Build an AI Data Analyst API

Turn plain English questions into SQL queries and charts with Gemini.

Introduction

30 Second Summary

Every business has databases full of answers. But getting those answers usually means writing SQL queries by hand or waiting for a data team to get back to you.

In this project, you will build an AI-powered API that converts natural language questions into SQL queries, runs them against a PostgreSQL database, and returns both the raw results and a ready-to-render chart.

What You'll Build

You'll type a question like "What are my top 5 best-selling products?" into the API, and watch it generate the SQL, pull the data from PostgreSQL, and hand back a bar chart you can render in any browser.

By the end of this project, you'll have:

  • A FastAPI backend that accepts plain English questions and returns executed SQL query results from a live PostgreSQL database.
  • An AI-powered text-to-SQL pipeline using Google Gemini that dynamically introspects your database schema to write accurate queries against any tables you point it at.
  • A safety validation layer that blocks destructive queries (DROP, DELETE, UPDATE) before they ever touch your data.
  • Auto-generated Plotly chart specifications returned alongside query results so any frontend can render interactive visualizations.
  • Secret Mission: Add conversation memory so follow-up questions like "now group that by month" actually work.

Are there any prerequisites?

This project assumes you have Python 3.10+ installed, PostgreSQL installed and running, and a Google AI Studio account with an API key.

If you don't have these yet, Step 1 will walk you through the setup.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Connect to PostgreSQL and Expose the Schema

Your database is seeded and your API skeleton is running. Now you need to give the AI a way to see what's inside that database.

For the AI to write accurate SQL, it needs to know your table names, column names, and data types. Without that context, it would hallucinate names that don't exist. In this step, you'll build a module that connects to PostgreSQL and dynamically reads the schema, then expose it through an API endpoint.

In this step, get ready to:
  • Create a database connection module using Psycopg 3.
  • Build a schema introspection function that reads table structure dynamically.
  • Add a /schema endpoint so you can verify the AI will have the right context.
Create a database connection module using Psycopg 3

Psycopg 3 is the PostgreSQL adapter for Python. You'll isolate the database connection logic in its own module so the rest of your app can import it cleanly.

  • Create a new file called database.py in your project folder (the same folder as main.py).
  • Add the following code to set up the connection:
import os
import psycopg
from dotenv import load_dotenv

load_dotenv()

DATABASE_URL = os.getenv("DATABASE_URL")


def get_connection():
    # Open a connection to PostgreSQL using the URL from .env
    return psycopg.connect(DATABASE_URL)

What does this code do?

  • load_dotenv() reads your .env file and loads the variables into the environment.
  • DATABASE_URL pulls the connection string you set up in Step 1.
  • get_connection() returns a live Psycopg 3 connection that any part of your app can use to run queries.
  • Save database.py.
  • Test the connection by running this command in your terminal:
python -c "from database import get_connection; conn = get_connection(); print('Connected!'); conn.close()"

You should see Connected! printed in your terminal.

Seeing a connection error?

  • Make sure PostgreSQL is running. If you stopped it, restart it before testing.
  • Check that your DATABASE_URL in .env matches your actual PostgreSQL username and password.
  • Confirm the ai_analyst database exists by running psql -l and looking for it in the list.

Help me debug my PostgreSQL connection error in Psycopg 3.

Build the schema introspection function

PostgreSQL stores metadata about every table and column in a system view called information_schema.columns. You can query it to discover your database structure without hardcoding any table names.

  • In database.py, add this function below get_connection():
def get_schema() -> str:
    # Query PostgreSQL's system catalog for table structure
    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute("""
                SELECT table_name, column_name, data_type
                FROM information_schema.columns
                WHERE table_schema = 'public'
                ORDER BY table_name, ordinal_position;
            """)
            rows = cur.fetchall()

    # Group columns by table name
    schema_dict = {}
    for table_name, column_name, data_type in rows:
        if table_name not in schema_dict:
            schema_dict[table_name] = []
        schema_dict[table_name].append(f"{column_name} ({data_type})")

    # Format into a readable string for the LLM
    schema_str = "Database Schema:\n"
    for table, columns in schema_dict.items():
        schema_str += f"\nTable: {table}\n"
        schema_str += f"  Columns: {', '.join(columns)}\n"

    return schema_str

What does this code do?

  • The SQL query hits information_schema.columns to pull every table name, column name, and data type from your database.
  • table_schema = 'public' filters to only your app's tables, skipping PostgreSQL's internal system tables.
  • The function groups columns by table, then formats everything into a plain-text string. This string is what gets fed to the AI as context later.
  • Save database.py.
  • Test the schema introspection by running this command:
python -c "from database import get_schema; print(get_schema())"

You should see a formatted list of all four tables (customers, order_items, orders, products) with their columns and data types printed in your terminal.

Seeing an empty schema or an error?

  • If the output is empty, you may not have run the seed script in Step 1. Run python seed.py to populate the database.
  • Make sure you added get_schema() below get_connection() in the same file. The function calls get_connection() internally, so it must be defined after it.

Help me debug why my schema introspection returns empty results.

Why dynamic introspection instead of hardcoding?

You could paste your table definitions directly into the AI prompt. But querying the schema dynamically means your API works with any database you point it at. Change the DATABASE_URL and the AI instantly knows the new structure.

✔️ Awesome, I've got everything!

Great. Make sure you saved database.py with both functions before moving on.

ⓧ I'd like to double check the full code

Your database.py should contain two sections, in this order:

  • The imports, environment loading, and the connection function (from the first chunk above).
  • The get_schema() function (from the second chunk above).

Scroll up and confirm each chunk is present in your file, in order.

Add a /schema endpoint to expose the introspected schema

Now you'll wire the schema function into your FastAPI app so you can verify it works through the browser.

  • In main.py, find this line at the top:
from fastapi import FastAPI
  • Add this import directly below it:
from database import get_schema
  • Now add the /schema endpoint at the bottom of main.py, after the health_check function:
@app.get("/schema")
def schema():
    return {"schema": get_schema()}

What does this endpoint do?

When someone hits GET /schema, FastAPI calls get_schema() from your database module. It reads the live table structure from PostgreSQL and returns it as JSON. Later, you'll feed this same string to the AI so it knows what columns exist.

  • Save main.py.

The FastAPI dev server auto-reloads when you save. Before you test, what do you think you'll see when you hit the /schema endpoint?

  • Open your browser and go to http://127.0.0.1:8000/docs.
  • Click the GET /schema row to expand it.
  • Click Try it out, then click Execute.

You should see a JSON response containing a formatted schema string listing all four tables (customers, order_items, orders, products) with their columns and data types.

Server not reloading or endpoint not showing?

  • If the Swagger UI doesn't show /schema, check your terminal for errors. A typo in the import line will crash the auto-reload.
  • Make sure database.py is in the same folder as main.py. Python can only import files from the same directory.
  • If you stopped the server, restart it by running fastapi dev in your terminal.

Help me fix my FastAPI /schema endpoint that isn't loading.

✔️ Awesome, I've got everything!

Great. Make sure you saved main.py before moving on.

ⓧ I'd like to double check the full code

Here is the full main.py at this point:

from fastapi import FastAPI
from database import get_schema

app = FastAPI(title="AI Data Analyst", version="1.0.0")


@app.get("/")
def health_check():
    return {"status": "healthy", "service": "AI Data Analyst"}


@app.get("/schema")
def schema():
    return {"schema": get_schema()}

Your API can now read and expose the database structure dynamically. Next up, you'll build the core feature: a chat endpoint where the AI turns plain English questions into SQL queries and runs them against your data.

Build the Chat-to-SQL Endpoint

Your /schema endpoint proves that FastAPI can read your database structure. Now it is time to build the feature that makes this project useful.

In this step, you will connect Google Gemini to your API so a user can type a plain English question, the AI generates the SQL, your app executes that SQL against PostgreSQL, and the results come back as structured JSON.

In this step, get ready to:
  • Set up the Gemini AI client to generate SQL from natural language questions.
  • Define Pydantic models to structure request and response data.
  • Create a POST /chat endpoint that orchestrates the full question-to-results flow.
Create the Gemini client and SQL generation prompt

The LLM is responsible for translating a plain English question into valid SQL. You will give it the database schema as context so it knows which tables and columns exist.

  • Create a new file called llm.py in your project folder (the same folder as main.py).
  • Add the module setup that initializes the Gemini client by pasting this code:
import os
from google import genai
from google.genai import types
from dotenv import load_dotenv

load_dotenv()

# Initialize the Gemini client using your API key from .env
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

MODEL = "gemini-3.5-flash"

What does this code do?

This block loads environment variables and creates a connection to the Gemini API. The genai.Client reads your GEMINI_API_KEY from the .env file you created in Step 1.

The MODEL constant stores the model ID so you only need to change it in one place if you swap models later.

  • Below the MODEL line, add the generate_sql function that sends questions to Gemini:
def generate_sql(question: str, schema: str) -> str:
    system_instruction = f"""You are a SQL expert. Given the following PostgreSQL database schema, 
generate a single SQL SELECT query to answer the user's question.

{schema}

Rules:
- Return ONLY the SQL query, no explanations, no markdown code fences
- Only generate SELECT queries (never INSERT, UPDATE, DELETE, DROP, etc.)
- Use PostgreSQL syntax
- If the question cannot be answered with the schema, return: SELECT 'Question cannot be answered with available data' AS error;
"""

    response = client.models.generate_content(
        model=MODEL,
        contents=question,
        config=types.GenerateContentConfig(
            system_instruction=system_instruction,
            temperature=0.0,
        ),
    )

    # Strip markdown code fences if Gemini wraps the SQL
    sql = response.text.strip()
    if sql.startswith("```"):
        lines = sql.split("\n")
        sql = "\n".join(lines[1:-1])

    return sql.strip()

What does this function do?

The system_instruction tells Gemini it is a SQL expert and feeds it your full database schema. The rules constrain Gemini to return only raw SQL with no extra text.

Setting temperature=0.0 makes the output deterministic. The same question always produces the same SQL.

The code fence stripping at the bottom handles cases where Gemini wraps the SQL in markdown backticks despite being told not to.

  • Save llm.py.
  • Verify the module loads correctly by running this in a new terminal:
python -c "from llm import generate_sql; print('LLM module loaded successfully')"

You should see LLM module loaded successfully printed in the terminal.

Seeing an import error?

  • Make sure your virtual environment is activated. You should see (venv) at the start of your terminal prompt.
  • Check that google-genai is installed. Run pip list | grep google-genai to confirm.
  • Confirm you are running the command from the same folder where llm.py lives.

Help me debug this import error in my llm.py module

✔️ Awesome, I've got everything!

Great. Double check you have saved llm.py before moving on.

ⓧ I'd like to double check the full code

Your llm.py has two sections. Scroll up and confirm both are present in this order:

  • Module setup: imports, load_dotenv(), client initialization, and MODEL constant.
  • The generate_sql function with system instruction, generate_content call, and code fence stripping.
Update main.py with request and response models

Before the endpoint can accept questions and return results, you need to define the data shapes. Pydantic models tell FastAPI exactly what fields to expect in the request body and what to include in the response.

  • Switch back to main.py in VS Code.
  • Find the import line at the top of the file:
from fastapi import FastAPI
from database import get_schema
  • Replace those two lines with these updated imports:
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from database import get_schema, get_connection
from llm import generate_sql

What changed in the imports?

  • HTTPException lets you return proper HTTP error codes to the client.
  • BaseModel from Pydantic gives you typed request and response validation.
  • get_connection is the database helper you built in Step 2 for running queries.
  • generate_sql is the function you just created in llm.py.
  • Find the line app = FastAPI(title="AI Data Analyst", version="1.0.0") and add the two Pydantic models directly below it:
class ChatRequest(BaseModel):
    question: str


class ChatResponse(BaseModel):
    question: str
    sql: str
    columns: list[str]
    rows: list[list]
    chart: dict | None = None

What do these models define?

ChatRequest expects a single field: the user's natural language question.

ChatResponse defines the shape of every response: the original question, the generated SQL, column names, data rows, and an optional chart (which you will add later).

  • Save main.py.
  • Check the terminal where fastapi dev is running. You should see it reload without errors.

Server crashed on reload?

  • If you see an import error mentioning llm, confirm llm.py is saved in the same folder as main.py.
  • If you see a syntax error, check that the ChatResponse model uses dict | None = None (Python 3.10+ syntax). Make sure your virtual environment uses Python 3.10 or higher.

Help me fix my main.py reload error

Wire it into a POST /chat endpoint

The models define the contract. Now you will build the endpoint function that orchestrates the full flow: receive a question, generate SQL, execute it, and return the structured results.

  • In main.py, scroll to the bottom of the file (after the /schema endpoint) and add the /chat endpoint:
@app.post("/chat", response_model=ChatResponse)
def chat(request: ChatRequest):
    schema_str = get_schema()

    sql = generate_sql(request.question, schema_str)

    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(sql)
            columns = [desc[0] for desc in cur.description]
            rows = [list(row) for row in cur.fetchall()]

    return ChatResponse(
        question=request.question,
        sql=sql,
        columns=columns,
        rows=rows,
        chart=None,
    )

What does this endpoint do?

  • get_schema() fetches the live database structure so Gemini always has up-to-date context.
  • generate_sql() sends the question and schema to Gemini, which returns a SQL query.
  • The with get_connection() block executes that SQL against PostgreSQL and extracts column names from cur.description.
  • chart=None is a placeholder. You will add chart generation in a later step.
  • Save main.py.
  • Confirm the dev server reloads successfully in your terminal.

Before you test, think about what you expect to happen. If you send the question "What are the top 5 products by total revenue?" to Gemini, do you think it can figure out the correct SQL from the schema alone?

  • Open your browser to http://127.0.0.1:8000/docs.
  • Find the POST /chat endpoint in the Swagger UI and click it to expand.
  • Click Try it out.
  • Replace the example JSON body with this:
{
  "question": "What are the top 5 products by total revenue?"
}
  • Click Execute.

You should see a 200 response with a JSON body containing the generated SQL, column names, and rows of data showing products ranked by revenue.

What should the response look like?

The exact SQL will vary slightly depending on Gemini's output, but you should see a SELECT query that joins products and order_items, calculates revenue (quantity times unit_price), groups by product, and limits to 5 results. The rows field contains the actual data pulled from your database.

Getting a 500 Internal Server Error?

  • Check the terminal running fastapi dev for the full error traceback.
  • If you see an error about the API key, make sure your .env file has a valid GEMINI_API_KEY value (not the placeholder your-api-key-here).
  • If you see a connection error, confirm your PostgreSQL server is still running and the ai_analyst database exists.

Help me debug my /chat endpoint 500 error

✔️ Awesome, I've got everything!

Great. Make sure both llm.py and main.py are saved before moving on.

ⓧ I'd like to double check the full code

Your main.py should have these sections in order. Scroll through and confirm each is present:

  • Four import lines: fastapi, pydantic, database, llm.
  • The app = FastAPI(...) line.
  • ChatRequest and ChatResponse Pydantic models.
  • The health_check endpoint at /.
  • The schema endpoint at /schema.
  • The chat endpoint at POST /chat that calls get_schema(), generate_sql(), executes the SQL, and returns a ChatResponse.

You just built the core AI feature. Your API accepts a plain English question and returns real data from your database. Next up, you will add a safety layer so that dangerous queries like "delete all orders" never reach the database.

Add SQL Safety Validation

Your /chat endpoint works. It takes a question, generates SQL, runs it, and returns results. But there is a serious problem lurking underneath.

Right now, nothing stops the AI from generating destructive queries like DELETE or DROP TABLE. In this step, you will build a safety layer that inspects every generated query and rejects anything that is not a read-only SELECT.

In this step, get ready to:
  • See the vulnerability in your current endpoint by sending a destructive prompt.
  • Build a SQL validator module that blocks non-SELECT queries.
  • Integrate the validator into your /chat endpoint and confirm dangerous prompts are rejected.
See the problem firsthand

Your current endpoint passes whatever SQL Gemini generates straight to the database. Let's see what happens when you ask something destructive.

  • Switch back to the Swagger UI at http://127.0.0.1:8000/docs in your browser.
  • Expand the POST /chat section and click Try it out.

Before you send this, what do you think will happen if you ask the AI to delete your data? Will it refuse, or will it try to do it?

  • Replace the request body with the following and click Execute:
{"question": "Delete all the orders"}

You will see a 500 Internal Server Error. Gemini generated a DELETE statement, and your code tried to execute it against the database. The 500 error happened because a DELETE query does not return columns for your code to read.

Your data was only saved by accident

The crash happened AFTER the DELETE was sent to PostgreSQL. Your data survived only because psycopg rolled back the transaction when the subsequent code error occurred.

A slightly different query structure (like a DROP TABLE) could bypass this accidental safety net entirely. You need an intentional validation layer that catches dangerous queries before they ever reach the database.

Build a SQL validator that rejects non-SELECT statements

The fix is a dedicated module that inspects every generated SQL string before execution. It will strip out comments, normalize whitespace, and reject anything that is not a pure SELECT query.

  • Create a new file called safety.py in your project folder (the same folder as main.py).
  • Add the following code to safety.py:
import re


# Keywords that indicate a destructive or modifying query
FORBIDDEN_KEYWORDS = [
    "INSERT", "UPDATE", "DELETE", "DROP",
    "ALTER", "TRUNCATE", "CREATE", "GRANT",
    "REVOKE", "EXEC", "EXECUTE",
]


def validate_sql(sql: str) -> tuple[bool, str]:
    # Strip single-line comments
    cleaned = re.sub(r"--.*$", "", sql, flags=re.MULTILINE)
    # Strip multi-line comments
    cleaned = re.sub(r"/\*.*?\*/", "", cleaned, flags=re.DOTALL)
    # Normalize whitespace and convert to uppercase for checking
    cleaned = " ".join(cleaned.split()).strip().upper()

    # Must start with SELECT
    if not cleaned.startswith("SELECT"):
        return False, "Only SELECT queries are allowed."

    # Scan for forbidden keywords
    for keyword in FORBIDDEN_KEYWORDS:
        pattern = r"\b" + keyword + r"\b"
        if re.search(pattern, cleaned):
            return False, f"Forbidden keyword detected: {keyword}"

    return True, "Query is safe."

What does this code do?

  • FORBIDDEN_KEYWORDS is a list of SQL keywords that modify or destroy data. Any query containing these words will be blocked.
  • The function first strips SQL comments (both -- and /* */ styles) so attackers cannot hide keywords inside comments.
  • It normalizes whitespace and uppercases the string so that tricks like extra spaces or mixed case do not bypass the check.
  • It returns a tuple: (True, "Query is safe.") if the query passes, or (False, "reason") if it fails.
  • Save safety.py.
  • Test the validator directly by running this command in your terminal:
python -c "from safety import validate_sql; print(validate_sql('DELETE FROM orders'))"

You should see (False, 'Only SELECT queries are allowed.') printed in your terminal. The validator caught the dangerous query.

Seeing a ModuleNotFoundError?

  • Make sure you saved the file as safety.py (not Safety.py or safety.txt) in the same folder as main.py.
  • Make sure your terminal is in the project folder. Run ls and confirm you see safety.py listed.
  • Confirm your virtual environment is still active (you should see the environment name in your terminal prompt).

Why can't Python find my safety module?

Integrate the validator into the /chat endpoint

The validator works in isolation. Now you need to wire it into your /chat endpoint so every generated query is checked before it touches the database.

  • In main.py, find the import section at the top of the file. Locate this line:
from llm import generate_sql
  • Add a new import line directly below it, so it looks like this:
from llm import generate_sql
from safety import validate_sql
  • Now scroll down to the chat function. Find this line:
    sql = generate_sql(request.question, schema_str)
  • Add the validation check directly below that line, before the with get_connection() block:
    # Validate the SQL before executing it
    is_safe, message = validate_sql(sql)
    if not is_safe:
        raise HTTPException(status_code=400, detail=message)

What does this addition do?

After Gemini generates SQL, validate_sql() inspects it. If the query is unsafe, the endpoint immediately returns a 400 error with a clear message. The query never reaches the database.

  • Save main.py. The dev server will auto-reload.

Time to test. Before you send the dangerous prompt again, what do you think will happen this time? Will you get a 500 error, a 400 error, or a successful response?

  • Return to the Swagger UI at http://127.0.0.1:8000/docs.
  • Expand POST /chat and click Try it out.
  • Send this request body:
{"question": "Delete all the orders"}

You should see a 400 response with a message like Only SELECT queries are allowed. instead of the 500 crash from before. Your safety layer caught it.

  • Try one more dangerous prompt to confirm. Send this request body:
{"question": "Drop the products table"}

You should see another 400 response. This time the message may say Forbidden keyword detected: DROP depending on what SQL Gemini generated.

  • Finally, confirm that legitimate queries still work. Send this:
{"question": "What are the top 5 products by total revenue?"}

You should see a 200 response with SQL, column names, and rows of product data. Safe queries pass through as normal.

Still seeing a 500 error on dangerous prompts?

  • Check that you added from safety import validate_sql at the top of main.py.
  • Verify the validation code is placed AFTER sql = generate_sql(...) and BEFORE the with get_connection() block.
  • Check your terminal for reload errors. If the server failed to restart, you may have a syntax error in main.py.

Why is my SQL validation not catching dangerous queries?

✔️ Awesome, I've got everything!

Great. Double check you have saved both safety.py and main.py.

ⓧ I'd like to double check the full code

Here is the full safety.py:

import re


FORBIDDEN_KEYWORDS = [
    "INSERT", "UPDATE", "DELETE", "DROP",
    "ALTER", "TRUNCATE", "CREATE", "GRANT",
    "REVOKE", "EXEC", "EXECUTE",
]


def validate_sql(sql: str) -> tuple[bool, str]:
    cleaned = re.sub(r"--.*$", "", sql, flags=re.MULTILINE)
    cleaned = re.sub(r"/\*.*?\*/", "", cleaned, flags=re.DOTALL)
    cleaned = " ".join(cleaned.split()).strip().upper()

    if not cleaned.startswith("SELECT"):
        return False, "Only SELECT queries are allowed."

    for keyword in FORBIDDEN_KEYWORDS:
        pattern = r"\b" + keyword + r"\b"
        if re.search(pattern, cleaned):
            return False, f"Forbidden keyword detected: {keyword}"

    return True, "Query is safe."

Here is the full main.py:

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from database import get_schema, get_connection
from llm import generate_sql
from safety import validate_sql

app = FastAPI(title="AI Data Analyst", version="1.0.0")


class ChatRequest(BaseModel):
    question: str


class ChatResponse(BaseModel):
    question: str
    sql: str
    columns: list[str]
    rows: list[list]
    chart: dict | None = None
@app.get("/")
def health_check():
    return {"status": "healthy", "service": "AI Data Analyst"}


@app.get("/schema")
def schema():
    return {"schema": get_schema()}


@app.post("/chat", response_model=ChatResponse)
def chat(request: ChatRequest):
    schema_str = get_schema()

    sql = generate_sql(request.question, schema_str)

    is_safe, message = validate_sql(sql)
    if not is_safe:
        raise HTTPException(status_code=400, detail=message)

    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(sql)
            columns = [desc[0] for desc in cur.description]
            rows = [list(row) for row in cur.fetchall()]

    return ChatResponse(
        question=request.question,
        sql=sql,
        columns=columns,
        rows=rows,
        chart=None,
    )

Your API now has a real safety layer. Destructive queries are caught and rejected with a clear error message before they ever reach the database. Next up, you will make the results more visual by auto-generating charts from query output.

Auto-Generate Charts from Query Results

Your API safely executes AI-generated SQL and hands back raw rows and columns. That works, but reading a table of numbers is hard.

In this step, you will build a chart generator that detects the shape of query results and returns a Plotly chart specification alongside the data. Any frontend can then render an interactive visualization from that JSON without extra work.

In this step, get ready to:
  • Build a chart generator that detects whether results are time-series, categorical, or numeric.
  • Wire chart generation into the /chat response so every query returns a ready-to-render Plotly figure.
Build a chart generator that detects data shape and picks a chart type

The chart generator inspects the first column of query results to decide which visualization fits best. Dates become line charts, string categories become bar charts, and everything else defaults to a bar chart.

  • Create a new file called charts.py in your project folder (the same folder where main.py lives).
  • Add the first part of the generate_chart function, which extracts data and detects the chart type:
def generate_chart(columns: list[str], rows: list[list]) -> dict | None:
    # Return None if data is too sparse to chart
    if not rows or len(columns) < 2:
        return None

    # Extract x and y values from the first two columns
    first_col_values = [row[0] for row in rows]
    x_label = columns[0]
    y_label = columns[1]
    x_values = [str(v) for v in first_col_values]
    y_values = [row[1] for row in rows]

    # Detect data shape to pick the right chart type
    is_date = _looks_like_dates(first_col_values)
    is_category = all(isinstance(v, str) for v in first_col_values)

    if is_date:
        chart_type = "scatter"
        mode = "lines+markers"
    elif is_category:
        chart_type = "bar"
        mode = None
    else:
        chart_type = "bar"
        mode = None

What does this code do?

The function takes column names and row data from a SQL result. It treats the first column as the x-axis and the second as the y-axis.

The _looks_like_dates call (which you will write next) checks whether x-axis values are dates. If they are, a line chart (scatter with lines+markers) makes sense. If the values are strings like product categories, a bar chart is more readable.

  • Save charts.py and confirm the new file exists by running:
ls charts.py

You should see charts.py listed in the output.

Don't see charts.py?

Make sure you created the file inside the same project folder where main.py lives. Run ls without arguments to see all files in your current directory.

help me find where charts.py was saved

  • Still in charts.py, add this code directly below the last line (continuing inside the generate_chart function):
    # Build the Plotly trace dict
    trace = {
        "x": x_values,
        "y": y_values,
        "type": chart_type,
        "name": y_label,
    }
    if mode:
        trace["mode"] = mode

    # Assemble the full Plotly figure
    chart = {
        "data": [trace],
        "layout": {
            "title": f"{y_label} by {x_label}",
            "xaxis": {"title": x_label},
            "yaxis": {"title": y_label},
        },
    }

    return chart

What does the trace and layout do?

A Plotly figure is a JSON object with two main keys: data (the plotted values) and layout (axis titles, chart title, formatting). The trace dict holds one series of x/y data points.

By returning this dict from the API, any frontend with plotly.js can render it directly. No image rendering happens on the server.

  • Below generate_chart (leave two blank lines after return chart), add this helper function that detects date values:
def _looks_like_dates(values) -> bool:
    import datetime

    # Check the first 5 values for date patterns
    for v in values[:5]:
        if isinstance(v, (datetime.date, datetime.datetime)):
            return True
        if isinstance(v, str):
            try:
                datetime.date.fromisoformat(v)
                return True
            except (ValueError, TypeError):
                pass
    return False

Why check only the first 5 values?

Checking every row would waste time on large result sets. Sampling the first 5 is enough to detect whether the column holds dates. PostgreSQL returns date columns as Python datetime.date objects through Psycopg 3, so the isinstance check catches them immediately.

  • Save charts.py.
  • Verify the module loads without errors by running:
python -c "from charts import generate_chart; print('charts module ready')"

You should see charts module ready printed with no errors.

Seeing a NameError or SyntaxError?

A NameError about _looks_like_dates means the helper function is missing or incorrectly indented. It must be a top-level function (no indentation) placed below generate_chart.

A SyntaxError usually means inconsistent indentation inside generate_chart. Confirm every line inside the function uses exactly 4 spaces.

help me debug my charts.py import error

✔️ Awesome, I've got everything!

Great. Double check you have saved charts.py and the import succeeds before moving on.

ⓧ I'd like to double check the full code

Your charts.py contains two functions. Scroll up and confirm you have each chunk in order:

  • generate_chart starting with the data extraction, then chart type detection, then trace building, chart assembly, and return chart.
  • _looks_like_dates as a separate top-level function below generate_chart, checking the first 5 values for date patterns.
Integrate chart generation into the /chat response

The chart module is ready. Now you need to call it from your /chat endpoint and include the result in the response.

  • In main.py, find this line near the top of your imports:
from safety import validate_sql
  • Add this import directly below it:
from charts import generate_chart

What does this import do?

This brings in the generate_chart function you just built. FastAPI will call it after every successful query to produce a chart specification.

  • Now find this return block at the end of the chat function:
    return ChatResponse(
        question=request.question,
        sql=sql,
        columns=columns,
        rows=rows,
        chart=None,
    )
  • Replace that entire return block with the following, which generates the chart before returning:
    chart = generate_chart(columns, rows)

    return ChatResponse(
        question=request.question,
        sql=sql,
        columns=columns,
        rows=rows,
        chart=chart,
    )

What changed?

Instead of always returning chart=None, the endpoint now calls generate_chart(columns, rows) with the query results. If the data has at least two columns, the response includes a full Plotly figure dict. If not, it gracefully returns None.

  • Save main.py. The FastAPI dev server auto-reloads.

Before you test, what do you expect to see in the response when you ask a time-based question? Think about what chart type the function would pick for monthly data.

  • Switch to your browser and return to the Swagger UI at http://127.0.0.1:8000/docs.
  • Expand the POST /chat endpoint and click Try it out.
  • Enter this request body:
{
  "question": "Show me monthly revenue for 2024"
}
  • Click Execute.

You should see a JSON response containing a chart field with a Plotly figure. The chart will have "type": "scatter" and "mode": "lines+markers" because the x-axis values are dates.

Not seeing a chart field?

If chart is null, the query might be returning fewer than 2 columns. Check the columns field in the response to verify at least two columns exist.

If you see a 500 error, check your terminal for the traceback. Common causes are a missing import in main.py or incorrect indentation in charts.py.

help me debug why my /chat endpoint is not returning a chart

Why Plotly JSON instead of an image?

Returning a Plotly figure as JSON is lightweight and frontend-agnostic. Any web page with plotly.js loaded can call Plotly.newPlot(element, data, layout) and get an interactive chart with hover, zoom, and pan.

✔️ Awesome, I've got everything!

Great. Double check you have saved both charts.py and main.py before continuing.

ⓧ I'd like to double check main.py

Your main.py now has two changes from the previous step. Scroll up and confirm:

  • The import section includes from charts import generate_chart after the safety import.
  • The chat function calls chart = generate_chart(columns, rows) before the return statement, and passes chart=chart instead of chart=None.

Your API now returns raw data AND a ready-to-render chart in a single response. Any question you ask gets back SQL, rows, and a visualization spec tailored to the data shape. Next up, you will extend this into a conversational experience with follow-up questions.

Secret mission

Add Conversation Memory

Your API answers questions one at a time, but real analysts ask follow-ups. "Now group that by month" or "show me just the top 3" only make sense when the AI remembers what came before. Add conversation memory so your API handles multi-turn questions like a real data analyst would.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • A local PostgreSQL database called ai_analyst with sample e-commerce data.
  • A project folder on your Desktop containing Python source files, a virtual environment, and a .env file.

Keep everything running

No action needed. Choose this if you are still actively building or experimenting with new questions and chart types.

  • Your FastAPI dev server continues running at http://127.0.0.1:8000.
  • The ai_analyst database persists locally until you explicitly drop it.
  • Your virtual environment stays active in the terminal session.

Pause - I'll come back to this later

Shut down running processes to free up memory. Your files and database remain intact for next time.

  • Stop the FastAPI dev server by pressing Ctrl+C in the terminal where it is running.
  • Deactivate your virtual environment by running:
deactivate
  • The ai_analyst database persists until you delete it. No action needed to pause it.
  • When you come back, reactivate your virtual environment and run fastapi dev again to resume.

Delete - I don't want to use this again

Remove all project resources and start fresh. This is permanent.

  • Stop the FastAPI dev server by pressing Ctrl+C if it is still running.
  • Deactivate your virtual environment by running:
deactivate
  • Drop the PostgreSQL database by running:
dropdb ai_analyst
  • Remove your project folder from the Desktop. On macOS or Linux, run:
cd ~/Desktop
rm -rf ai-data-analyst
  • On Windows, run:
cd ~\Desktop
rmdir /s /q ai-data-analyst
  • That's everything. Your machine is back to its original state.

Nice Work!

Nice Work!

You just built a fully functional AI Data Analyst API from scratch. You can now type plain English questions and get back SQL, data, and ready-to-render charts from your own PostgreSQL database.

You've learned how to:

  • Build an AI-powered text-to-SQL pipeline using Google Gemini that introspects your database schema and converts natural language questions into executable queries.
  • Add a SQL safety validation layer that blocks destructive queries like DROP, DELETE, and UPDATE before they ever reach the database.
  • Auto-generate Plotly chart specifications based on query result shapes so any frontend can render interactive visualizations from your API response.
  • Secret Mission: Added conversation memory so follow-up questions like "now group that by month" resolve correctly using prior context.

Ready to quiz yourself?