Visualization agent with Plotly

The Visualization Agent uses the LLM to generate Plotly code dynamically. It then executes this code and returns the chart as JSON for the UI.

Chart generation pipeline

How the visualization agent generates charts from data

text2sql_agent.py
python
def viz_agent(state: AgentState) -> AgentState:
    """Generate a Plotly visualization from query results"""
    query_result = state["query_result"]
    graph_type = state["graph_type"]
    question = state["question"]

    try:
        results = json.loads(query_result)
        df = pd.DataFrame(results)

        prompt = f"""Generate Python code using Plotly to visualize this data.

Question: {question}
Graph Type: {graph_type}
Columns: {df.columns.tolist()}
Sample Data: {df.head(5).to_dict('records')}

Requirements:
1. Use plotly.graph_objects or plotly.express
2. Data is already loaded as 'df' (pandas DataFrame)
3. Create a {graph_type} chart
4. Limit to top 20 rows if needed
5. Variable must be named 'fig'
6. Return ONLY Python code, no markdown
7. Do NOT include imports or fig.show()
"""

        response = completion(
            model=DEFAULT_MODEL,
            messages=[{"role": "system", "content": AGENT_CONFIGS["viz_agent"]["system_prompt"]},
                      {"role": "user", "content": prompt}],
            temperature=0.3
        )

        plotly_code = response.choices[0].message.content.strip()
        plotly_code = plotly_code.replace("```python", "").replace("```", "")

        # Execute the generated code
        exec_globals = {'df': df, 'pd': pd, 'go': go, 'px': px}
        exec(plotly_code, exec_globals)

        fig = exec_globals.get('fig')
        state["graph_json"] = fig.to_json() if fig else ""

    except Exception as e:
        print(f"Viz error: {e}")
        state["graph_json"] = ""

    return state

The Viz Agent generates and executes Plotly code dynamically.

Good security instinct! Yes, exec() is risky. We mitigate it by: 1) sandboxing the globals so the code can only access df, pandas, and plotly, 2) wrapping everything in try/except, and 3) not exposing file system or network access. In production, you would add further sandboxing or use a code execution service like E2B.

Validation checklist: Visualization agent concepts

Loading practice…

The visualization agent can now generate charts from any query result. Next, we will learn common Plotly patterns and best practices for polished charts.