First pass at getting layouts

lumen/ai/models.py

@@ -156,6 +156,7 @@ class VegaLiteSpecUpdate(BaseModel):
         Respect your step's scope; don't override previous steps."""
     )
 
+
 class LineChange(BaseModel):
     line_no: int = Field(description="The line number in the original text that needs to be changed.")
     replacement: str = Field(description="""
@@ -363,6 +364,94 @@ class DbtslQueryParams(BaseModel):
     )
 
 
+class VegaLitePlotSpec(BaseModel):
+    """Specification for a single plot to be generated independently."""
+
+    slug: str = Field(
+        description="""
+        Unique identifier for this plot (e.g., 'revenue_trend', 'top_5_categories').
+        Must be unique within the layout and descriptive of what's shown.
+        """
+    )
+
+    instruction: str = Field(
+        description="""
+        Detailed natural language instructions for generating this plot.
+        Include: mark type (bar/line/point/etc), x-axis field and type,
+        y-axis field and type, any color encoding, aggregations, filters,
+        and styling preferences. Be specific enough that each plot can be
+        generated independently without additional context.
+
+        Example: "Create a line chart showing revenue over time. X-axis:
+        'month' (temporal), Y-axis: 'revenue' (quantitative, aggregated as sum).
+        Use blue color for the line. Add points to mark each data point."
+        """
+    )
+
+    title: str | None = Field(
+        default=None,
+        description="Optional title for this subplot. If omitted, will be generated during plot creation."
+    )
+
+
+class VegaLiteRow(BaseModel):
+    """A row in the layout, containing one or more plots arranged horizontally."""
+
+    plot_slugs: list[str] = Field(
+        description="""
+        List of plot slugs for this row, in left-to-right order.
+        All plots in a row are arranged side-by-side (hconcat).
+        """
+    )
+
+
+class VegaLiteLayoutPlan(BaseModel):
+    """
+    Plan for creating a multi-plot visualization.
+
+    Layout structure:
+    - Single row with 1 plot → Simple plot (no concat)
+    - Single row with N plots → hconcat (side-by-side)
+    - Multiple rows with 1 plot each → vconcat (stacked)
+    - Multiple rows with varying plots → vconcat of hconcats (grid)
+    """
+
+    chain_of_thought: str = Field(
+        description="""
+        Explain your visualization strategy:
+        - What story are you telling with this layout?
+        - Why did you choose these specific plots?
+        - Why this arrangement (side-by-side vs stacked)?
+        - How do the plots complement each other?
+        """
+    )
+
+    overall_title: str | None = Field(
+        default=None,
+        description="Optional overall title for the entire visualization dashboard."
+    )
+
+    plots: list[VegaLitePlotSpec] = Field(
+        description="""
+        All plots to generate, listed in the order they should appear.
+        Each plot will be generated independently in parallel.
+        """
+    )
+
+    rows: list[VegaLiteRow] = Field(
+        description="""
+        Layout structure defining how plots are arranged.
+        Each row contains plot slugs that will be arranged horizontally.
+        Multiple rows are stacked vertically.
+
+        Examples:
+        - Side-by-side: [{"plot_slugs": ["plot1", "plot2", "plot3"]}]
+        - Stacked: [{"plot_slugs": ["plot1"]}, {"plot_slugs": ["plot2"]}]
+        - Grid: [{"plot_slugs": ["plot1", "plot2"]}, {"plot_slugs": ["plot3"]}]
+        """
+    )
+
+
 class QueryCompletionValidation(PartialBaseModel):
     """Validation of whether the executed plan answered the user's query"""
 

lumen/ai/prompts/LumenBaseAgent/retry_output.jinja2

@@ -26,11 +26,8 @@ Available schema context:
 Language: {{ language }}
 {% endif %}
 
-Original code with line numbers:
+Original code:
 ```{% if language %}{{ language }}{% endif %}
-{{ numbered_text }}
+{{ original_text }}
 ```
-
-Please restate this feedback in your own words, and then try to fix the code based on the feedback provided:
-{{ feedback }}
 {% endblock %}

lumen/ai/prompts/Planner/main.jinja2

@@ -8,6 +8,12 @@ You are the team lead responsible for creating a step-by-step plan to address us
 - If ❌ BLOCKED, find actors with "Provides" matching the missing Requires and add them as prior steps
 - Dependencies must be resolved in the correct order - providers before consumers
 
+🔄 CRITICAL: Handling Follow-up Requests for Additional Visualizations
+- When the user asks to "add" or "include" additional metrics/dimensions to an existing visualization:
+  * **For side-by-side comparisons** (e.g., "add Silver medals too", "also show Profit"): The SQL query MUST return ALL metrics in a SINGLE result set
+  * Instruct SQLAgent to MODIFY the existing query to include BOTH the original AND new metrics
+  * Do NOT create separate queries for each metric - they need to be in the same table for multi-plot layouts
+
 Ground Rules:
 - Respect dependency chains: assign tasks only when input Requires are met
 - Leverage existing memory instead of regenerating information if possible

lumen/ai/prompts/SQLAgent/main.jinja2

@@ -5,12 +5,19 @@
 Write a SQL query for the user's data transformation request, focusing on intent over literal interpretation.
 
 ## 🎯 PRIMARY RULES:
-1. Reuse materialized data - Check "Current Knowledge" first before writing CTEs
-2. Column quotes - Use double quotes: `"column_name"`
-3. String quotes - Single quotes: `'string_value'`
-4. No CREATE statements - System handles materialization
-5. No LIMIT clauses unless explicitly requested by the user - Pagination is automatically handled
-6. Progressive approach - Each step builds on previous
+1. **Reuse and extend** - When asked to "add" or "include" additional metrics:
+   - Check "Previous" query in Current Knowledge
+   - MODIFY the existing query to include ALL metrics in a SINGLE result set
+   - Example: If previous query counts gold medals, and user asks for silver too, return BOTH gold_count AND silver_count columns
+   - Use CASE statements, multiple aggregations, or pivoting to include all requested metrics
+2. Reuse materialized data - Check "Current Knowledge" first before writing CTEs
+3. Column quotes - Use double quotes: `"column_name"`
+4. String quotes - Single quotes: `'string_value'`
+5. No CREATE statements - System handles materialization
+6. No LIMIT clauses unless explicitly requested by the user - Pagination is automatically handled
+7. Progressive approach - Each step builds on previous
+
+Then EXTEND the previous query to return multiple columns, not create a separate query.
 
 ## Query Patterns:
 - NULL handling: Exclude unless specified

lumen/ai/prompts/VegaLiteAgent/annotate_plot.jinja2

@@ -1,80 +1,81 @@
-{% extends 'Actor/main.jinja2' %}
+{% extends 'LumenBaseAgent/retry_output.jinja2' %}
 
 {% block instructions %}
-Add strategic visual annotations to highlight key insights in your data.
+Add strategic visual annotations to highlight key insights.
 
-# Annotation Patterns
+# LAYERS
 
-**Reference lines**: Horizontal/vertical rules for meaningful thresholds
-- Use for goals, benchmarks, or natural breakpoints (e.g., "5,000+ turbines = major wind state")
-- Position with `y: {datum: value}` or `x: {datum: value}`
-- Use subtle styling: `opacity: 0.5`, muted colors, thin lines
-- Anchor labels to the line, not floating positions
+Wrap chart in `layer` array, then add annotation marks (rules, text, highlights). All layers share data/scales unless overridden.
 
-**Value labels**: Text on specific data points
-- Use only for outliers or endpoints that need emphasis
-- Apply `transform` with `filter` to select points
-- Don't label every bar - let the axis provide that information
+# PATTERNS
+
+**Reference lines**: Horizontal/vertical rules for thresholds (goals, benchmarks)
+- Position: `y: {datum: value}` or `x: {datum: value}`
+- Style: subtle colors, `opacity: 0.5`, `strokeDash: [4,4]`
+- Anchor labels to line
+
+**Value labels**: Text on outliers/endpoints only
+- Filter: `transform: [{filter: "datum.field > threshold"}]`
+- Don't label every point—let axes handle that
 
 **Conditional highlights**: Emphasize data meeting criteria
 - Layer filtered data with distinct color
-- Use `transform: [{filter: "datum.field > threshold"}]`
-- Make sure the threshold is meaningful, not arbitrary
-{% endblock %}
-
-{% block context %}
-Build off the following Vega-Lite yaml:
-```yaml
-{{ vega_spec }}
-```
+- Use meaningful thresholds
 {% endblock %}
 
 {% block examples %}
 # EXAMPLES
 
-Reference line with subtle styling:
+## Converting to layers (reference line):
+
+Before:
+```yaml
+mark: bar
+encoding:
+  x: {field: category, type: ordinal}
+  y: {field: value, type: quantitative}
+```
+
+After:
 ```yaml
 layer:
   - mark: bar
     encoding:
       x: {field: category, type: ordinal}
       y: {field: value, type: quantitative}
-  - mark: {type: rule, color: '#666666', size: 1, strokeDash: [4,4], opacity: 0.5}
-    encoding:
-      y: {datum: 5000}
-  - mark: {type: text, align: left, dx: 5, dy: -5, fontSize: 10, color: '#666666'}
+  - mark: {type: rule, color: '#666', size: 1, strokeDash: [4,4], opacity: 0.5}
+    encoding: {y: {datum: 5000}}
+  - mark: {type: text, align: left, dx: 5, dy: -5, fontSize: 10, color: '#666'}
     encoding:
       text: {value: "5,000"}
       y: {datum: 5000}
       x: {datum: 0}
 ```
 
-Value labels with filter:
+## Value labels with filter:
 ```yaml
 layer:
   - mark: line
     encoding:
       x: {field: year, type: ordinal}
       y: {field: value, type: quantitative}
   - mark: {type: text, align: left, dx: 5, fontSize: 11}
-    transform:
-      - filter: "datum.year == 2022"  # Only label final point
+    transform: [{filter: "datum.year == 2022"}]
     encoding:
       x: {field: year, type: ordinal}
       y: {field: value, type: quantitative}
       text: {field: value, type: quantitative, format: ".0f"}
 ```
 
-Conditional highlighting:
+## Conditional highlighting:
 ```yaml
 layer:
   - mark: bar
     encoding:
       x: {field: category, type: nominal}
       y: {field: value, type: quantitative}
   - mark: {type: bar, color: "#e74c3c"}
-    transform:
-      - filter: "datum.value > 500"
+    transform: [{filter: "datum.value > 500"}]
     encoding:
       x: {field: category, type: nominal}
       y: {field: value, type: quantitative}