Switch to Altair charts for dynamic visualizations

ai_data_analyst/sub_agents/visualization_agent.py

@@ -12,26 +12,26 @@
 logger.info("Visualization Agent module loaded.")
 
 
-def generate_chart(
-        chart_type: str,
+def validate_and_pass_chart_spec(
+        chart_type: str, # Retained for potential future use or logging, though not strictly needed for validation
         chart_spec: str,
 ) -> str:
     """
-    Generates an Altair chart specification based on the provided chart specification.
-    This tool receives a complete Vega-Lite chart specification and validates it.
+    Validates a chart specification (expected to be Altair/Vega-Lite JSON) and returns it if valid.
+    This tool receives a complete chart specification and validates it.
 
     Args:
-        chart_type: The type of chart (e.g., 'altair', 'bar', 'line', etc.)
-        chart_spec: A JSON string containing the complete Vega-Lite chart specification
+        chart_type: The type of chart (e.g., 'altair', 'bar', 'line', etc.). Helps in logging.
+        chart_spec: A JSON string containing the complete chart specification.
 
     Returns:
-        The validated chart specification as a JSON string, or an error message
+        The validated chart specification as a JSON string, or an error message.
     """
     logger.debug(
-        f"generate_chart called with chart_type: {chart_type}, chart_spec length: {len(chart_spec) if chart_spec else 0}")
+        f"validate_and_pass_chart_spec called with chart_type: {chart_type}, chart_spec length: {len(chart_spec) if chart_spec else 0}")
 
     if not chart_spec:
-        logger.error("No chart specification provided.")
+        logger.error("No chart specification provided to validate_and_pass_chart_spec.")
         return "Error: No chart specification provided."
 
     try:
@@ -62,10 +62,10 @@ def generate_chart(
 
 VISUALIZATION_AGENT_INSTRUCTION = """
 # ROLE AND GOAL
-You are a specialized "Consultative Visualization Agent." Your core function is to act as an expert data visualization designer. Your primary goal is to convert datasets into clear, compelling, and accurate visualizations by generating markdown that includes both Vega-Lite chart specifications and expert recommendations. You are also an expert consultant, expected to recommend the most effective visualization type to achieve a given analytical goal.
+You are a specialized "Consultative Visualization Agent." Your core function is to act as an expert data visualization designer. Your primary goal is to convert datasets into clear, compelling, and accurate visualizations by generating markdown that includes **Altair-compatible JSON chart specifications** and expert recommendations. You are also an expert consultant, expected to recommend the most effective visualization type to achieve a given analytical goal. Your JSON output for charts should be directly usable by tools that render Altair or Vega-Lite specifications (e.g., Streamlit's `st.altair_chart`).
 
 # CORE CAPABILITIES
-1.  **Vega-Lite Chart Generation:** You can create a wide variety of charts by generating their complete Vega-Lite JSON specification. Your capabilities include, but are not limited to: bar charts, histograms, scatter plots, box plots, line charts, and heatmaps.
+1.  **Altair/Vega-Lite Chart Generation:** You can create a wide variety of charts by generating their complete Altair-compatible JSON specification (which is effectively a Vega-Lite JSON). Your capabilities include, but are not limited to: bar charts, histograms, scatter plots, box plots, line charts, and heatmaps. Utilize Altair's expressive grammar to create effective visualizations.
 2.  **Visualization Recommendation:** Given a dataset and an analytical goal, you will determine the most effective chart type. For instance, you will recommend a bar chart over a pie chart for comparing many categories, or a box plot instead of a simple bar chart of means to show a distribution's spread.
 3.  **Aesthetic Best Practices:** You must ensure all generated chart specifications include clear and appropriate titles, axis labels, and legends to make them immediately understandable.
 
@@ -79,57 +79,60 @@ def generate_chart(
 # OUTPUT FORMAT
 Your output **MUST** be a single markdown code block with the language identifier `json`.
 This JSON object must contain two keys:
-1.  `"chart_spec"`: The value will be the complete Vega-Lite JSON specification for the chart.
+1.  `"chart_spec"`: The value will be the complete **Altair-compatible (Vega-Lite) JSON specification** for the chart. This JSON should be directly renderable by tools like Streamlit's `st.altair_chart`.
 2.  `"filterable_columns"`: The value will be a JSON array of strings, where each string is the name of a column from the dataset that is a good candidate for user-based filtering (e.g., categorical columns with a reasonable number of unique values).
 
 ### Example Output Structure
 ```json
 {
   "chart_spec": {
     "$schema": "https://vega.github.io/schema/vega-lite/v5.json",
-    "description": "Average Value by Category.",
+    "description": "Average Value by Category, showing regional breakdown.",
     "data": {
       "values": [
         {"category": "A", "region": "North", "mean_value": 521.3},
-        {"category": "B", "region": "South", "mean_value": 615.8}
+        {"category": "B", "region": "South", "mean_value": 615.8},
+        {"category": "A", "region": "East", "mean_value": 488.0},
+        {"category": "B", "region": "West", "mean_value": 701.2}
       ]
     },
     "mark": "bar",
     "encoding": {
       "x": {"field": "category", "type": "nominal", "title": "Category"},
-      "y": {"field": "mean_value", "type": "quantitative", "title": "Average Value"}
+      "y": {"field": "mean_value", "type": "quantitative", "title": "Average Value"},
+      "color": {"field": "region", "type": "nominal", "title": "Region"}
     },
-    "title": "Average Value by Category"
+    "title": "Average Value by Category and Region"
   },
   "filterable_columns": ["region", "category"]
 }
 ```
 
 # ACCESSIBILITY REQUIREMENTS
-- Use colorblind-friendly palettes
-- Include text alternatives for visual elements
-- Ensure sufficient color contrast
+- Use colorblind-friendly palettes.
+- Include text alternatives for visual elements (e.g., `description` field in the spec).
+- Ensure sufficient color contrast.
 
 # ENEM-OPTIMIZED VISUALIZATIONS
-- Score distribution histograms with Brazilian education context
-- Geographic heat maps for regional performance
-- Comparative bar charts for public vs private schools
+- Score distribution histograms with Brazilian education context.
+- Geographic heat maps for regional performance.
+- Comparative bar charts for public vs private schools.
 
 # KEY PRINCIPLES & CONSTRAINTS
-1.  **MARKDOWN FORMAT IS MANDATORY:** All outputs must be properly formatted markdown with a recommendation section followed by a `vega-lite` code block containing the chart specification.
-2.  **VEGA-LITE SPECIFICATION:** The chart specification within the code block **MUST** be valid Vega-Lite JSON. Do not generate code for other plotting libraries (e.g., Matplotlib, Plotly), SVG, or image files.
+1.  **MARKDOWN JSON FORMAT IS MANDATORY:** All outputs must be properly formatted markdown with a `json` code block containing the chart specification and filterable columns.
+2.  **ALTAIR/VEGA-LITE JSON SPECIFICATION:** The chart specification within the code block **MUST** be valid Altair-compatible (Vega-Lite) JSON. Do not generate code for other plotting libraries (e.g., Matplotlib, Plotly), SVG, or image files.
 3.  **NO DATA ANALYSIS:** You **DO NOT** perform calculations, aggregations, or statistical analysis. You visualize the data exactly as it is given to you. If the data needs to be aggregated for the chart, you must rely on the input `dataset` to already be in the correct aggregated format.
 4.  **CLARITY IS KING:** Prioritize creating visualizations that are easy to understand. Avoid misleading representations like truncated y-axes or overly complex charts.
-5.  **CONCISE RECOMMENDATIONS:** Keep recommendations focused and practical, explaining the visualization choice without interpreting the data's meaning or implications.
+5.  **CONCISE OUTPUT:** Keep the JSON output focused on the chart specification and filterable columns as described.
 
 """
 # Create the agent instance
 visualization_agent = LlmAgent(
     name="visualization_agent_tool",
     model="gemini-2.5-flash",
     instruction=VISUALIZATION_AGENT_INSTRUCTION,
-    description="Generates specifications for data visualizations by calling the `generate_chart` tool based on provided datasets and user requests.",
-    tools=[generate_chart], # Provide the agent with the tool it can use
+    description="Generates specifications for data visualizations by calling the `validate_and_pass_chart_spec` tool based on provided datasets and user requests.",
+    tools=[validate_and_pass_chart_spec], # Provide the agent with the tool it can use
     output_key="visualization_agent_output_key",
     generate_content_config=types.GenerateContentConfig(
         temperature=0.1,

streamlit_app.py

@@ -84,11 +84,11 @@ def display_message_content(message_text):
                 if 'data' in spec_for_render:
                     del spec_for_render['data']
 
-                st.vega_lite_chart(filtered_data, spec_for_render, use_container_width=True)
+                st.altair_chart(filtered_data, use_container_width=True, **spec_for_render)
             else:
                 # For non-filterable charts (no filterable_columns) or charts
                 # with data from a URL (data is empty), render the spec directly.
-                st.vega_lite_chart(chart_spec, use_container_width=True)
+                st.altair_chart(chart_spec, use_container_width=True)
 
         except json.JSONDecodeError as e:
             st.error(f"Failed to parse chart JSON: {e}")