fix: Update chart generation and LLM agent functionality

app.py

@@ -16,7 +16,9 @@ logging.getLogger('PIL').setLevel(logging.WARNING)
 
 app = Flask(__name__, static_folder=os.path.join(os.path.dirname(__file__), '..', 'static'))
 
-CORS(app)
+# Configure CORS to allow all origins for development
+CORS(app, origins=["*"], supports_credentials=True)
+
 agent = LLM_Agent()
 
 UPLOAD_FOLDER = os.path.join(os.path.dirname(__file__), '..', 'data', 'uploads')
@@ -53,7 +55,12 @@ def plot():
 @app.route('/static/<path:filename>')
 def serve_static(filename):
     logging.info(f"Serving static file: {filename}")
-    return send_from_directory(app.static_folder, filename)
+    response = send_from_directory(app.static_folder, filename)
+    # Add CORS headers for images
+    response.headers.add('Access-Control-Allow-Origin', '*')
+    response.headers.add('Access-Control-Allow-Headers', 'Content-Type')
+    response.headers.add('Access-Control-Allow-Methods', 'GET')
+    return response
 
 @app.route('/upload', methods=['POST'])
 def upload_file():

chart_generator.py

@@ -27,32 +27,54 @@ class ChartGenerator:
                 missing_cols.append(y)
         if missing_cols:
             logging.error(f"Missing columns in data: {missing_cols}")
+            logging.info(f"Available columns: {list(self.data.columns)}")
             raise ValueError(f"Missing columns in data: {missing_cols}")
         
-        fig, ax = plt.subplots()
+        # Clear any existing plots
+        plt.clf()
+        plt.close('all')
+        
+        fig, ax = plt.subplots(figsize=(10, 6))
+        
         for y in y_cols:
             color = plot_args.get('color', None)
             if plot_args.get('chart_type', 'line') == 'bar':
                 ax.bar(self.data[x_col], self.data[y], label=y, color=color)
             else:
-                ax.plot(self.data[x_col], self.data[y], label=y, color=color)
+                ax.plot(self.data[x_col], self.data[y], label=y, color=color, marker='o')
         
         ax.set_xlabel(x_col)
+        ax.set_ylabel('Value')
+        ax.set_title(f'{plot_args.get("chart_type", "line").title()} Chart')
         ax.legend()
+        ax.grid(True, alpha=0.3)
         
+        # Rotate x-axis labels if needed
+        if len(self.data[x_col]) > 5:
+            plt.xticks(rotation=45)
 
         chart_filename = 'chart.png'
         output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'static', 'images')
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
+            logging.info(f"Created output directory: {output_dir}")
         
         full_path = os.path.join(output_dir, chart_filename)
         
         if os.path.exists(full_path):
             os.remove(full_path)
+            logging.info(f"Removed existing chart file: {full_path}")
         
-        plt.savefig(full_path)
+        # Save with high DPI for better quality
+        plt.savefig(full_path, dpi=300, bbox_inches='tight', facecolor='white')
+        plt.close(fig)
         
-        logging.info(f"Chart generated and saved to {full_path}")
+        # Verify file was created
+        if os.path.exists(full_path):
+            file_size = os.path.getsize(full_path)
+            logging.info(f"Chart generated and saved to {full_path} (size: {file_size} bytes)")
+        else:
+            logging.error(f"Failed to create chart file at {full_path}")
+            raise FileNotFoundError(f"Chart file was not created at {full_path}")
         
         return os.path.join('static', 'images', chart_filename)

llm_agent.py

@@ -10,6 +10,7 @@ import os
 from dotenv import load_dotenv
 import ast
 import requests
+import json
 
 load_dotenv()
 
@@ -36,7 +37,7 @@ class LLM_Agent:
     def process_request(self, data):
         start_time = time.time()
         logging.info(f"Processing request data: {data}")
-        query = data['query']
+        query = data.get('query', '')
         data_path = data.get('file_path')
         model_choice = data.get('model', 'bart')
 
@@ -49,8 +50,16 @@ class LLM_Agent:
             else:
                 logging.info(f"File exists at path: {data_path}")
 
-        # Few-shot + persona prompt for Flan-UL2 (best model)
-        flan_prompt = (
+        # Re-initialize data processor and chart generator if a file is specified
+        if data_path:
+            self.data_processor = DataProcessor(data_path)
+            # Log loaded columns
+            loaded_columns = self.data_processor.get_columns()
+            logging.info(f"Loaded columns from data: {loaded_columns}")
+            self.chart_generator = ChartGenerator(self.data_processor.data)
+
+        # Enhanced prompt for better model responses
+        enhanced_prompt = (
             "You are VizBot, an expert data visualization assistant. "
             "Given a user's natural language request about plotting data, output ONLY a valid Python dictionary with keys: x, y, chart_type, and color (if specified). "
             "Do not include any explanation or extra text.\n\n"
@@ -66,73 +75,94 @@ class LLM_Agent:
             f"User: {query}\nOutput:"
         )
 
-        # Re-initialize data processor and chart generator if a file is specified
-        if data_path:
-            self.data_processor = DataProcessor(data_path)
-            # Log loaded columns
-            loaded_columns = self.data_processor.get_columns()
-            logging.info(f"Loaded columns from data: {loaded_columns}")
-            self.chart_generator = ChartGenerator(self.data_processor.data)
-
-        if model_choice == 'bart':
-            # Use local fine-tuned BART model
-            inputs = self.query_tokenizer(query, return_tensors="pt", max_length=512, truncation=True)
-            outputs = self.query_model.generate(**inputs, max_length=100, num_return_sequences=1)
-            response_text = self.query_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        elif model_choice == 'flan-t5-base':
-            # Use Hugging Face Inference API with Flan-T5-Base model
-            api_url = "https://api-inference.huggingface.co/models/google/flan-t5-base"
-            headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACEHUB_API_TOKEN')}", "Content-Type": "application/json"}
-            response = requests.post(api_url, headers=headers, json={"inputs": flan_prompt})
-            if response.status_code != 200:
-                logging.error(f"Hugging Face API error: {response.status_code} {response.text}")
-                response_text = "Error: Unable to get response from Flan-T5-Base API. Please try again later."
-            else:
-                try:
-                    resp_json = response.json()
-                    response_text = resp_json[0]['generated_text'] if isinstance(resp_json, list) else resp_json.get('generated_text', '')
-                except Exception as e:
-                    logging.error(f"Error parsing Hugging Face API response: {e}, raw: {response.text}")
-                    response_text = f"Error: Unexpected response from Flan-T5-Base API."
-        elif model_choice == 'flan-ul2':
-            # Use Hugging Face Inference API with Flan-UL2 model
-            api_url = "https://api-inference.huggingface.co/models/google/flan-ul2"
-            # Corrected model name to "google/flan-ul2" does not exist, use "google/flan-t5-xxl" as best available
-            api_url = "https://api-inference.huggingface.co/models/google/flan-t5-xxl"
-            headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACEHUB_API_TOKEN')}", "Content-Type": "application/json"}
-            response = requests.post(api_url, headers=headers, json={"inputs": flan_prompt})
-            if response.status_code != 200:
-                logging.error(f"Hugging Face API error: {response.status_code} {response.text}")
-                response_text = "Error: Unable to get response from Flan-T5-XXL API. Please try again later."
+        try:
+            if model_choice == 'bart':
+                # Use local fine-tuned BART model
+                inputs = self.query_tokenizer(query, return_tensors="pt", max_length=512, truncation=True)
+                outputs = self.query_model.generate(**inputs, max_length=100, num_return_sequences=1)
+                response_text = self.query_tokenizer.decode(outputs[0], skip_special_tokens=True)
+            elif model_choice == 'flan-t5-base':
+                # Use Hugging Face Inference API with Flan-T5-Base model
+                api_url = "https://api-inference.huggingface.co/models/google/flan-t5-base"
+                headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACEHUB_API_TOKEN')}"}
+                payload = {"inputs": enhanced_prompt}
+                
+                response = requests.post(api_url, headers=headers, json=payload, timeout=30)
+                if response.status_code != 200:
+                    logging.error(f"Hugging Face API error: {response.status_code} {response.text}")
+                    response_text = "{'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}"
+                else:
+                    try:
+                        resp_json = response.json()
+                        response_text = resp_json[0]['generated_text'] if isinstance(resp_json, list) else resp_json.get('generated_text', '')
+                        if not response_text:
+                            response_text = "{'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}"
+                    except Exception as e:
+                        logging.error(f"Error parsing Hugging Face API response: {e}, raw: {response.text}")
+                        response_text = "{'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}"
+            elif model_choice == 'flan-ul2':
+                # Use Hugging Face Inference API with Flan-T5-XXL model (best available)
+                api_url = "https://api-inference.huggingface.co/models/google/flan-t5-xxl"
+                headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACEHUB_API_TOKEN')}"}
+                payload = {"inputs": enhanced_prompt}
+                
+                response = requests.post(api_url, headers=headers, json=payload, timeout=30)
+                if response.status_code != 200:
+                    logging.error(f"Hugging Face API error: {response.status_code} {response.text}")
+                    response_text = "{'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}"
+                else:
+                    try:
+                        resp_json = response.json()
+                        response_text = resp_json[0]['generated_text'] if isinstance(resp_json, list) else resp_json.get('generated_text', '')
+                        if not response_text:
+                            response_text = "{'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}"
+                    except Exception as e:
+                        logging.error(f"Error parsing Hugging Face API response: {e}, raw: {response.text}")
+                        response_text = "{'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}"
             else:
-                try:
-                    resp_json = response.json()
-                    response_text = resp_json[0]['generated_text'] if isinstance(resp_json, list) else resp_json.get('generated_text', '')
-                except Exception as e:
-                    logging.error(f"Error parsing Hugging Face API response: {e}, raw: {response.text}")
-                    response_text = f"Error: Unexpected response from Flan-T5-XXL API."
-        else:
-            # Default fallback to local fine-tuned BART model
-            inputs = self.query_tokenizer(query, return_tensors="pt", max_length=512, truncation=True)
-            outputs = self.query_model.generate(**inputs, max_length=100, num_return_sequences=1)
-            response_text = self.query_tokenizer.decode(outputs[0], skip_special_tokens=True)
+                # Default fallback to local fine-tuned BART model
+                inputs = self.query_tokenizer(query, return_tensors="pt", max_length=512, truncation=True)
+                outputs = self.query_model.generate(**inputs, max_length=100, num_return_sequences=1)
+                response_text = self.query_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        logging.info(f"LLM response text: {response_text}")
-        try:
-            plot_args = ast.literal_eval(response_text)
-        except (SyntaxError, ValueError):
-            plot_args = {'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}
-            logging.warning(f"Invalid LLM response. Using default plot args: {plot_args}")
-        if LLM_Agent.validate_plot_args(plot_args):
+            logging.info(f"LLM response text: {response_text}")
+            
+            # Clean and parse the response
+            response_text = response_text.strip()
+            if response_text.startswith("```") and response_text.endswith("```"):
+                response_text = response_text[3:-3].strip()
+            if response_text.startswith("python"):
+                response_text = response_text[6:].strip()
+            
+            try:
+                plot_args = ast.literal_eval(response_text)
+            except (SyntaxError, ValueError) as e:
+                logging.warning(f"Invalid LLM response: {e}. Response: {response_text}")
+                plot_args = {'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}
+            
+            if not LLM_Agent.validate_plot_args(plot_args):
+                logging.warning("Invalid plot arguments. Using default.")
+                plot_args = {'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'}
+            
             chart_path = self.chart_generator.generate_chart(plot_args)
-        else:
-            logging.warning("Invalid plot arguments. Using default.")
-            chart_path = self.chart_generator.generate_chart({'x': 'Year', 'y': ['Sales'], 'chart_type': 'line'})
-        verified = self.image_verifier.verify(chart_path, query)
-        end_time = time.time()
-        logging.info(f"Processed request in {end_time - start_time} seconds")
-        return {
-            "response": response_text,
-            "chart_path": chart_path,
-            "verified": verified
-        }
+            verified = self.image_verifier.verify(chart_path, query)
+            
+            end_time = time.time()
+            logging.info(f"Processed request in {end_time - start_time} seconds")
+            
+            return {
+                "response": response_text,
+                "chart_path": chart_path,
+                "verified": verified
+            }
+            
+        except Exception as e:
+            logging.error(f"Error processing request: {e}")
+            end_time = time.time()
+            logging.info(f"Processed request in {end_time - start_time} seconds")
+            
+            return {
+                "response": f"Error: {str(e)}",
+                "chart_path": "",
+                "verified": False
+            }