added together ai agent

controller.py

@@ -367,7 +367,7 @@ async def csv_chat(request: Dict, authorization: str = Header(None)):
         # groq_answer = await asyncio.to_thread(groq_chat, decoded_url, query)
         # logger.info("groq_answer:", groq_answer)
         
-        result = await query_csv_agent(decoded_url, query)
+        result = await query_csv_agent(decoded_url, query, chat_id)
         logger.info("together ai csv answer == >", result)
         return {"answer": result}
      
@@ -633,7 +633,7 @@ async def csv_chart(request: dict, authorization: str = Header(None)):
         #     return {"image_url": image_public_url}
             # return FileResponse(groq_result, media_type="image/png")
             
-        result = await query_csv_agent(csv_url,query)
+        result = await query_csv_agent(csv_url, query, chat_id)
         logger.info("together ai result ==>", result)
         return {"orchestrator_response": jsonable_encoder(result)}
           

python_code_executor_service.py

@@ -1,3 +1,6 @@
+import os
+from supabase import create_client, Client
+from dotenv import load_dotenv
 import uuid
 import matplotlib.pyplot as plt
 from pathlib import Path
@@ -14,6 +17,10 @@ import seaborn as sns
 import scipy.stats as stats
 from pydantic import BaseModel
 
+from supabase_service import upload_file_to_supabase
+
+# Load environment variables from .env file
+load_dotenv()
 
 class CodeResponse(BaseModel):
     """Container for code-related responses"""
@@ -136,126 +143,107 @@ class PythonExecutor:
             'plots': plots
         }
     
-    def save_plot_dummy(self, plot_data: bytes, description: str) -> str:
+    async def save_plot_to_supabase(self, plot_data: bytes, description: str, chat_id: str) -> str:
         """
-        Save plot to charts folder and return a dummy URL
+        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: Dummy URL for the chart
+            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 (even though we're using dummy URLs, we still save it)
+        # Save the plot locally first
         with open(filepath, 'wb') as f:
             f.write(plot_data)
             
-        # Return a dummy URL
-        return f"https://example.com/charts/{filename}"
+        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 _looks_like_structured_data(self, output: str) -> bool:
+        """Helper to detect JSON-like or array-like output"""
+        output = output.strip()
+        return (
+            output.startswith('{') and output.endswith('}') or  # JSON object
+            output.startswith('[') and output.endswith(']') or  # Array
+            '\n' in output and '=' in output  # Python console output
+        )
     
-    # def process_response(self, response: CsvChatResult) -> str:
-    #     """
-    #     Process the CsvChatResult response and generate formatted output
+    async def process_response(self, response: CsvChatResult, chat_id: str) -> str:
+        """
+        Process the CsvChatResult response and generate formatted output
+        with markdown code blocks for structured data.
         
-    #     Args:
-    #         response (CsvChatResult): Response from CSV analysis
+        Args:
+            response (CsvChatResult): Response from CSV analysis
+            chat_id (str): ID of the chat session
             
-    #     Returns:
-    #         str: Formatted output with results and dummy image URLs
-    #     """
-    #     output_parts = []
+        Returns:
+            str: Formatted output with results and image URLs
+        """
+        output_parts = []
         
-    #     # Add casual response
-    #     output_parts.append(response.casual_response)
+        # Add casual response
+        output_parts.append(response.casual_response)
         
-    #     # Process analysis operations
-    #     for operation in response.analysis_operations:
-    #         # Execute the code
-    #         result = self.execute_code(operation.code.code)
+        # Process analysis operations
+        for operation in response.analysis_operations:
+            # Execute the code
+            result = self.execute_code(operation.code.code)
             
-    #         # Add operation description
-    #         output_parts.append(f"\n{operation.description}:")
+            # Add operation description
+            output_parts.append(f"\n{operation.description}:")
             
-    #         # Add output or error
-    #         if result['error']:
-    #             output_parts.append(f"Error: {result['error']['message']}")
-    #         else:
-    #             output_parts.append(result['output'].strip())
-        
-    #     # Process charts if they exist
-    #     if response.charts:
-    #         output_parts.append("\nVisualizations:")
-            
-    #         for chart in response.charts:
-    #             if chart.code:
-    #                 # Execute the chart code
-    #                 result = self.execute_code(chart.code)
-                    
-    #                 if result['plots']:
-    #                     # Save each generated plot and get dummy URL
-    #                     for plot_data in result['plots']:
-    #                         dummy_url = self.save_plot_dummy(plot_data, chart.image_description)
-    #                         output_parts.append(f"\n{chart.image_description}")
-    #                         output_parts.append(f"![{chart.image_description}]({dummy_url})")
-    #                 elif result['error']:
-    #                     output_parts.append(f"\nError generating {chart.image_description}: {result['error']['message']}")
-        
-    #     return "\n".join(output_parts)
-    
-    def process_response(self, response: CsvChatResult) -> str:
-     """
-     Process the CsvChatResult response and generate formatted output
-     with markdown code blocks for structured data.
-     """
-     output_parts = []
-    
-     # Add casual response
-     output_parts.append(response.casual_response)
-    
- # Process analysis operations
-     for operation in response.analysis_operations:
-         # Execute the code
-         result = self.execute_code(operation.code.code)
-        
-         # Add operation description
-         output_parts.append(f"\n{operation.description}:")
-        
-         # Add output or error with markdown wrapping
-         if result['error']:
-             output_parts.append("```python\n" + f"Error: {result['error']['message']}" + "\n```")
-         else:
-            output = result['output'].strip()
-            if self._looks_like_structured_data(output):  # New helper method
-                output_parts.append("```python\n" + output + "\n```")
+            # Add output or error with markdown wrapping
+            if result['error']:
+                output_parts.append("```python\n" + f"Error: {result['error']['message']}" + "\n```")
             else:
-                output_parts.append(output)
-    
-     # Process charts remains the same
-     if response.charts:
-        output_parts.append("\nVisualizations:")
-        for chart in response.charts:
-            if chart.code:
-                result = self.execute_code(chart.code)
-                if result['plots']:
-                    for plot_data in result['plots']:
-                        dummy_url = self.save_plot_dummy(plot_data, chart.image_description)
-                        output_parts.append(f"\n{chart.image_description}")
-                        output_parts.append(f"![{chart.image_description}]({dummy_url})")
-                elif result['error']:
-                    output_parts.append("```python\n" + f"Error generating {chart.image_description}: {result['error']['message']}" + "\n```")
-    
-     return "\n".join(output_parts)
-
-    def _looks_like_structured_data(self, output: str) -> bool:
-     """Helper to detect JSON-like or array-like output"""
-     output = output.strip()
-     return (
-        output.startswith('{') and output.endswith('}') or  # JSON object
-        output.startswith('[') and output.endswith(']') or  # Array
-        '\n' in output and '=' in output  # Python console output
-     )
+                output = result['output'].strip()
+                if self._looks_like_structured_data(output):
+                    output_parts.append("```python\n" + output + "\n```")
+                else:
+                    output_parts.append(output)
+        
+        # Process charts
+        if response.charts:
+            output_parts.append("\nVisualizations:")
+            for chart in response.charts:
+                if chart.code:
+                    result = self.execute_code(chart.code)
+                    if result['plots']:
+                        for plot_data in 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"\nError uploading chart: {str(e)}")
+                    elif result['error']:
+                        output_parts.append("```python\n" + f"Error generating {chart.image_description}: {result['error']['message']}" + "\n```")
+        
+        return "\n".join(output_parts)

together_ai_llama_agent.py

@@ -113,7 +113,7 @@ def create_csv_agent(df: pd.DataFrame, max_retries: int = 1) -> Agent:
     raise RuntimeError(f"Failed to create agent after {max_retries} attempts")
 
 
-async def query_csv_agent(csv_url: str, question: str) -> str:
+async def query_csv_agent(csv_url: str, question: str, chat_id: str) -> str:
     """Query the CSV agent with a DataFrame and question and return formatted output"""
     
     # Get the DataFrame from the CSV URL
@@ -140,6 +140,6 @@ async def query_csv_agent(csv_url: str, question: str) -> str:
         print("Chat Result Original Object:", chat_result)
     
     # Process and format the response
-    formatted_output = executor.process_response(chat_result)
+    formatted_output = executor.process_response(chat_result, chat_id)
     
     return formatted_output