Chart Improvement Agent - Context-Based Iterative Optimization

In data visualization scenarios, users often need to adjust charts multiple times to achieve ideal results. The traditional approach is to regenerate the entire chart, but this loses previous context and user intent. AskTable's ChartImprovementAgent adopts a more intelligent approach: preserve original context, understand improvement intent, iteratively optimize charts.

Problem Background

When users generate charts with AI, common improvement requests include:

• •Style Adjustment: "Change bar chart to line chart"
• •Data Filtering: "Show only top 10"
• •Color Modification: "Use blue color scheme"
• •Layout Optimization: "Put legend at bottom"
• •Title Modification: "Change title to '2024 Sales Trend'"

If regenerating each time, LLM might:

• •Lose original data source reference
• •Change chart's core logic
• •Ignore previous optimizations

ChartImprovementAgent Architecture

Core Design

1. Context Preservation

ChartImprovementAgent receives complete context information during initialization:

class ChartImprovementAgent:
    def __init__(
        self,
        parent_nodes_context: list[dict],  # Parent node data
        original_question: str,             # Original question
        original_code: str,                 # Original JSX code
        improvement_request: str,             # Improvement request
    ):
        self.parent_nodes_context = parent_nodes_context
        self.original_question = original_question
        self.original_code = original_code
        self.improvement_request = improvement_request

Parent node context includes:

• •Node ID
• •SQL query
• •Data description
• •DataFrame structure (column names, sample data)

parent_nodes_context = [
    {
        "id": "node_123",
        "sql": "SELECT region, SUM(sales) as total FROM orders GROUP BY region",
        "description": "Sales summary by region",
        "dataframe": {
            "columns": ["region", "total"],
            "sample_data": [
                {"region": "East China", "total": 1000000},
                {"region": "North China", "total": 800000},
            ]
        }
    }
]

2. System Prompt Construction

Format context information into system prompt:

# Format parent node information
parent_info_parts = []
for i, node in enumerate(parent_nodes_context, 1):
    parent_info_parts.append(f"Node {i} (ID: {node['id']}):")
    if node.get("description"):
        parent_info_parts.append(f"  Description: {node['description']}")
    if node.get("sql"):
        parent_info_parts.append(f"  SQL: {node['sql']}")
    if node.get("dataframe"):
        df = node["dataframe"]
        parent_info_parts.append(f"  Columns: {', '.join(df.get('columns', []))}")
        if df.get("sample_data"):
            parent_info_parts.append(f"  Sample rows: {len(df['sample_data'])}")
    parent_info_parts.append("")

formatted_parent_info = "\n".join(parent_info_parts)

# Build system Prompt
system_prompt = get_prompt("agent/canvas/edit_chart").compile(
    formatted_parent_info=formatted_parent_info,
    original_question=original_question,
    original_code=original_code,
)

3. Improvement Tool

Agent provides submit_improved_chart tool for submitting improved code:

def submit_improved_chart(
    self,
    question: str = Field(
        ...,
        description="The rewritten question that naturally incorporates both original requirement and the improvement",
    ),
    description: str = Field(
        ...,
        description="Brief description of the improvements made (1-2 sentences), or error reason if code is None",
    ),
    code: str | None = Field(
        None,
        description="The improved JSX code for the chart component. Set to None if improvement failed.",
    ),
) -> str:
    if code is not None:
        try:
            # Compile JSX code
            self.compiled_code = compile_jsx(code)
            self.source_code = code
            log.info("Improved chart code compiled successfully")
        except Exception as e:
            log.error(f"Failed to compile improved chart code: {str(e)}")
            raise ValueError(f"Code compilation failed: {str(e)}")
    else:
        # Improvement failed, record reason
        self.compiled_code = None
        self.source_code = None
        log.warning(f"Chart improvement failed: {description}")

    self.submitted_description = description
    self.submitted_question = question

    status_msg = "error" if code is None else "success"
    return json.dumps({
        "question": question,
        "description": description,
        "status": status_msg,
        "has_code": code is not None,
    })

4. Question Rewriting

Agent needs to merge original question and improvement request into new question description:

Example:

• •Original question: "Show sales by region"
• •Improvement request: "Change to line chart and sort by sales descending"
• •Rewritten: "Show sales by region, use line chart, sort by sales descending"

This maintains question completeness for subsequent understanding and maintenance.

Practical Application Examples

Example 1: Chart Type Conversion

Original Code:

<code>
import { BarChart } from '@/components/charts';

function RegionSales() {
  const data = load_dataframe('df_region_sales');

  return (
    <BarChart
      data={data}
      xField="region"
      yField="total"
      title="Sales by Region"
    />
  );
}
</code>

Improvement Request: "Change to line chart"

Improved Code:

<code>
import { LineChart } from '@/components/charts';

function RegionSales() {
  const data = load_dataframe('df_region_sales');

  return (
    <LineChart
      data={data}
      xField="region"
      yField="total"
      title="Sales by Region"
      smooth={true}
    />
  );
}
</code>

Example 2: Data Filtering

Improvement Request: "Show only top 5 regions by sales"

Improved Code:

<code>
import { BarChart } from '@/components/charts';

function RegionSales() {
  const data = load_dataframe('df_region_sales');

  // Sort by sales descending, take top 5
  const topData = data
    .sort((a, b) => b.total - a.total)
    .slice(0, 5);

  return (
    <BarChart
      data={topData}
      xField="region"
      yField="total"
      title="Top 5 Regions by Sales"
    />
  );
}
</code>

Example 3: Style Optimization

Improvement Request: "Use blue color scheme, put legend at bottom"

Improved Code:

<code>
import { BarChart } from '@/components/charts';

function RegionSales() {
  const data = load_dataframe('df_region_sales');

  return (
    <BarChart
      data={data}
      xField="region"
      yField="total"
      title="Sales by Region"
      color="#1890ff"
      legend={{
        position: 'bottom'
      }}
    />
  );
}
</code>

Example 4: Complex Improvement

Improvement Request: "Add average line, annotate highest value, use gradient colors"

Improved Code:

<code>
import { BarChart } from '@/components/charts';

function RegionSales() {
  const data = load_dataframe('df_region_sales');

  // Calculate average
  const average = data.reduce((sum, item) => sum + item.total, 0) / data.length;

  // Find highest value
  const maxItem = data.reduce((max, item) => item.total > max.total ? item : max);

  return (
    <div>
      <BarChart
        data={data}
        xField="region"
        yField="total"
        title="Sales by Region"
        color={{
          type: 'gradient',
          colors: ['#1890ff', '#52c41a']
        }}
        annotations={[
          {
            type: 'line',
            start: ['min', average],
            end: ['max', average],
            style: { stroke: '#ff4d4f', lineDash: [4, 4] },
            text: { content: `Average: ${average.toFixed(0)}`, position: 'end' }
          },
          {
            type: 'text',
            position: [maxItem.region, maxItem.total],
            content: `Highest: ${maxItem.total}`,
            style: { fill: '#ff4d4f', fontWeight: 'bold' }
          }
        ]}
      />
    </div>
  );
}
</code>

Error Handling

1. Compilation Failure

If improved code cannot be compiled, Agent returns error information:

try:
    self.compiled_code = compile_jsx(code)
except Exception as e:
    raise ValueError(f"Code compilation failed: {str(e)}")

2. Improvement Not Feasible

If LLM determines improvement request cannot be implemented, can return code=None:

def submit_improved_chart(
    self,
    question: str,
    description: str,
    code: str | None = None,  # None means improvement failed
):
    if code is None:
        # Record failure reason
        self.compiled_code = None
        self.source_code = None
        log.warning(f"Chart improvement failed: {description}")

Example:

• •Improvement request: "Add 3D effect"
• •LLM response: code=None, description="Current chart library does not support 3D effects"

3. Data Source Verification

Ensure improved code still references correct data sources:

# Extract load_dataframe references
load_df_pattern = r"load_dataframe\(\s*['\"]( df_[A-Za-z0-9]+)['\"]\s*\)"
referenced_dataframes = re.findall(load_df_pattern, code)

# Verify data sources exist
missing_ids = set(referenced_dataframes) - set(self.data_workspace.keys())
if missing_ids:
    raise ValueError(f"Referenced dataframes {missing_ids} are not in the data workspace")

Comparison with ChartNodeAgent

User Experience Optimization

1. Progressive Improvement

Users can iteratively improve charts multiple times:

User: Show sales by region
→ Generate bar chart

User: Change to line chart
→ Improved to line chart

User: Show only top 5
→ Added data filtering

User: Use blue color scheme
→ Adjusted color scheme

2. Context Preservation

Each improvement retains all previous optimizations:

// First improvement: change to line chart
<LineChart ... />

// Second improvement: show only top 5
const topData = data.slice(0, 5);
<LineChart data={topData} ... />

// Third improvement: use blue color scheme
const topData = data.slice(0, 5);
<LineChart data={topData} color="#1890ff" ... />

3. Failure Degradation

If improvement fails, preserve original chart:

def get_result(self) -> dict:
    if self.source_code is None:
        return {
            "code": None,
            "compiled_code": None,
            "question": self.submitted_question or self.original_question,
            "description": self.submitted_description,
            "status": "error",
            "error": self.submitted_description,
        }

    return {
        "code": self.source_code,
        "compiled_code": self.compiled_code,
        "question": self.submitted_question,
        "description": self.submitted_description,
        "status": "success",
        "error": None,
    }

Performance Optimization

1. Incremental Compilation

Only compile improved code, not entire project:

self.compiled_code = compile_jsx(code)  # Single file compilation

2. Cache Reuse

Reuse parent node data, avoid repeated queries:

# Parent node data already includes DataFrame
parent_nodes_context = [
    {
        "id": "node_123",
        "dataframe": cached_dataframe  # Reuse cache
    }
]

3. Parallel Processing

Multiple improvement requests can be processed in parallel:

# Process multiple improvements in parallel
agents = [
    ChartImprovementAgent(..., improvement_request="Change to line chart"),
    ChartImprovementAgent(..., improvement_request="Show only top 5"),
]

results = await asyncio.gather(*[agent.run() for agent in agents])

Summary

AskTable's ChartImprovementAgent achieves intelligent chart iterative optimization through the following technologies:

1. •Context Preservation: Save original question, code, and data sources
2. •Intent Understanding: LLM understands user's improvement needs
3. •Question Rewriting: Merge original question and improvement request
4. •Code Verification: JSX compilation and data source verification
5. •Error Degradation: Preserve original chart when improvement fails

This design not only improves user experience but also ensures the reliability and consistency of chart improvement, being an important component of the AI-driven data visualization system.

Related Resources

• •AskTable Canvas Documentation
• •React Chart Library