fix: chart_generator

chart_generator.py

@@ -16,15 +16,28 @@ class ChartGenerator:
         start_time = time.time()
         logging.info(f"Generating chart with arguments: {plot_args}")
         
+        # Validate columns before plotting
+        x_col = plot_args['x']
+        y_cols = plot_args['y']
+        missing_cols = []
+        if x_col not in self.data.columns:
+            missing_cols.append(x_col)
+        for y in y_cols:
+            if y not in self.data.columns:
+                missing_cols.append(y)
+        if missing_cols:
+            logging.error(f"Missing columns in data: {missing_cols}")
+            raise ValueError(f"Missing columns in data: {missing_cols}")
+        
         fig, ax = plt.subplots()
-        for y in plot_args['y']:
+        for y in y_cols:
             color = plot_args.get('color', None)
             if plot_args.get('chart_type', 'line') == 'bar':
-                ax.bar(self.data[plot_args['x']], self.data[y], label=y, color=color)
+                ax.bar(self.data[x_col], self.data[y], label=y, color=color)
             else:
-                ax.plot(self.data[plot_args['x']], self.data[y], label=y, color=color)
+                ax.plot(self.data[x_col], self.data[y], label=y, color=color)
         
-        ax.set_xlabel(plot_args['x'])
+        ax.set_xlabel(x_col)
         ax.legend()