update files

agents/sql_agent/graph.py

@@ -37,8 +37,8 @@ def build_graph(visualize: bool = True) -> StateGraph:
         "detect_off_topic",
         lambda state: state['error'], 
         path_map={
-            # True: "generate_answer",  
-            True: "get_db_info",
+            True: "generate_answer",  
+            # True: "get_db_info",
             False: "get_db_info"
         }
     )

agents/sql_agent/nodes.py

@@ -332,6 +332,180 @@ def render_visualization(state: SQLAgentState) -> SQLAgentState:
     return state
 
 
+def render_visualization(state: SQLAgentState) -> SQLAgentState:
+    """
+    Render the visualization from formatted data.
+    Output: path to saved image file.
+    """
+    import matplotlib.pyplot as plt
+    import os
+    import uuid
+    from typing import Dict, Any, Optional
+
+    def save_fig(fig: plt.Figure) -> str:
+        """Save figure to file and return the file path."""
+        try:
+            output_dir = "output/plots"
+            os.makedirs(output_dir, exist_ok=True)
+            file_path = os.path.join(output_dir, f"visualization_{uuid.uuid4().hex[:8]}.png")
+            fig.savefig(file_path, format="png", bbox_inches="tight", dpi=100)
+            plt.close(fig)
+            return file_path
+        except Exception as e:
+            print(f"Error saving figure: {e}")
+            return ""
+
+    def validate_data(data: Dict[str, Any], required_keys: list) -> bool:
+        """Validate that data contains all required keys and has valid values."""
+        if not all(key in data for key in required_keys):
+            return False
+        # Check if there's actual data to plot
+        if "values" in data and not data["values"]:
+            return False
+        if "yValues" in data and not data["yValues"]:
+            return False
+        return True
+
+    def render_line(data: Dict[str, Any]) -> Optional[str]:
+        """Render line chart."""
+        required_keys = ["xValues", "yValues"]
+        if not validate_data(data, required_keys):
+            return None
+            
+        try:
+            fig, ax = plt.subplots(figsize=(10, 6))
+            x = data["xValues"]
+            for series in data["yValues"]:
+                if len(x) == len(series["data"]):
+                    ax.plot(x, series["data"], label=series.get("label", ""), marker='o')
+            
+            ax.set_xlabel(data.get("xAxisLabel", "X"))
+            ax.set_ylabel(data.get("yAxisLabel", "Y"))
+            ax.set_title(data.get("title", ""))
+            if any(series.get("label") for series in data["yValues"]):
+                ax.legend()
+            plt.tight_layout()
+            return save_fig(fig)
+        except Exception as e:
+            print(f"Error rendering line chart: {e}")
+            return None
+
+    def render_bar(data: Dict[str, Any], horizontal: bool = False) -> Optional[str]:
+        """Render bar chart (vertical or horizontal)."""
+        required_keys = ["labels", "values"]
+        if not validate_data(data, required_keys) or not data["values"]:
+            return None
+
+        try:
+            fig, ax = plt.subplots(figsize=(10, 6))
+            labels = data["labels"]
+            n_series = len(data["values"])
+            width = 0.8 / max(1, n_series)  # Prevent division by zero
+            x_indexes = list(range(len(labels)))
+            
+            for i, series in enumerate(data["values"]):
+                if not series["data"]:  # Skip empty series
+                    continue
+                    
+                offset = (i - n_series / 2) * width + width / 2
+                if horizontal:
+                    ax.barh(
+                        [x + offset for x in x_indexes],
+                        series["data"],
+                        height=width,
+                        label=series.get("label", f"Series {i+1}")
+                    )
+                    ax.set_yticks(x_indexes)
+                    ax.set_yticklabels(labels)
+                    ax.set_xlabel(data.get("xAxisLabel", "Value"))
+                    ax.set_ylabel(data.get("yAxisLabel", "Category"))
+                else:
+                    ax.bar(
+                        [x + offset for x in x_indexes],
+                        series["data"],
+                        width=width,
+                        label=series.get("label", f"Series {i+1}")
+                    )
+                    ax.set_xticks(x_indexes)
+                    ax.set_xticklabels(labels, rotation=45, ha='right')
+                    ax.set_xlabel(data.get("xAxisLabel", "Category"))
+                    ax.set_ylabel(data.get("yAxisLabel", "Value"))
+            
+            if any(series.get("label") for series in data["values"]):
+                ax.legend()
+            ax.set_title(data.get("title", ""))
+            plt.tight_layout()
+            return save_fig(fig)
+        except Exception as e:
+            print(f"Error rendering {'horizontal ' if horizontal else ''}bar chart: {e}")
+            return None
+
+    def render_scatter(data: Dict[str, Any]) -> Optional[str]:
+        """Render scatter plot."""
+        required_keys = ["series"]
+        if not validate_data(data, required_keys):
+            return None
+
+        try:
+            fig, ax = plt.subplots(figsize=(10, 6))
+            for series in data["series"]:
+                if not series.get("data"):
+                    continue
+                xs = [point.get("x", 0) for point in series["data"]]
+                ys = [point.get("y", 0) for point in series["data"]]
+                if len(xs) == len(ys):
+                    ax.scatter(
+                        xs, 
+                        ys, 
+                        label=series.get("label"),
+                        alpha=0.6,
+                        edgecolors='w'
+                    )
+            
+            ax.set_xlabel(data.get("xAxisLabel", "X"))
+            ax.set_ylabel(data.get("yAxisLabel", "Y"))
+            ax.set_title(data.get("title", ""))
+            if any(series.get("label") for series in data["series"]):
+                ax.legend()
+            plt.tight_layout()
+            return save_fig(fig)
+        except Exception as e:
+            print(f"Error rendering scatter plot: {e}")
+            return None
+
+    # Main function logic
+    data = state.get("formatted_data_for_visualization")
+    visualization = state.get("visualization", "none")
+    state["visualization_output"] = None
+
+    if not data or visualization == "none":
+        return state
+
+    try:
+        renderers = {
+            "line": lambda: render_line(data),
+            "bar": lambda: render_bar(data, horizontal=False),
+            "horizontal_bar": lambda: render_bar(data, horizontal=True),
+            "scatter": lambda: render_scatter(data)
+        }
+
+        if visualization in renderers:
+            image_path = renderers[visualization]()
+            if image_path and os.path.exists(image_path):
+                state["visualization_output"] = image_path
+            else:
+                state["error"] = "Failed to generate visualization: No valid data to display"
+        else:
+            state["error"] = f"Unsupported visualization type: {visualization}"
+
+    except Exception as e:
+        state["error"] = f"Error in visualization: {str(e)}"
+        print(f"Visualization error: {e}")
+
+    state["step"] = "render_visualization"
+    return state
+
+
 def finalize_output(state: SQLAgentState) -> SQLAgentState:
     """
     Node hợp nhất kết quả cuối cùng (answer, visualization_output, error, ...).
@@ -372,7 +546,8 @@ def detect_off_topic(state: SQLAgentState) -> SQLAgentState:
     )
     metadata = {
         "topic": "Database Queries",
-        "additional_context": "The database is about ecommerce products with tables: products, laptops, phones, tablets, promotions, category"
+        "additional_context": "Only accept queries related to the data on Database/CSV"
+        # "additional_context": "The database is about ecommerce products with tables: products, laptops, phones, tablets, promotions, category"
     }
 
     validation_result = validator.validate(question, metadata=metadata)
@@ -461,7 +636,8 @@ def generate_answer(state: SQLAgentState) -> SQLAgentState:
         state['error'] = state['error'] or "No results found."
         if state["off_topic"] == "OFF_TOPIC":
             state['error'] = "The question is off-topic." 
-            state["answer"] = "Sorry, I can't assist you with that request."  
+            # state["answer"] = "Sorry, I can't assist you with that request."
+            state["answer"] = "Sorry, I can only help you with questions about the data! What information would you like to explore from the data?"
     state['step'] = 'generate_answer'
     return state
 

app.py

@@ -464,11 +464,17 @@ else:
                                         st.write(step_details.get('answer', ''))
                                         chat_history.append({"role": "assistant", "content": step_details.get('answer', ''), "timestamp": datetime.now()})
                                     elif step_name == 'render_visualization':
-                                        # with st.expander("Chart", expanded=False):
-                                        st.image(step_details.get('visualization_output', ''))
+                                        try:
+                                            visualization_output = step_details.get('visualization_output')
+                                            if visualization_output and os.path.exists(visualization_output):
+                                                st.image(visualization_output)
+                                            else:
+                                                print("No visualization was generated for this query.")
+                                        except Exception as e:
+                                            print(f"Could not display visualization: {str(e)}")
                             except Exception as e:
                                 err = f"SQL Agent Error: {e}"
-                                st.error(err)
+                                print(err)
                                 chat_history.append({"role": "assistant", "content": err, "timestamp": datetime.now()})
                     else:
                         # Use DataFrame agent for selected CSV