> ## Documentation Index
> Fetch the complete documentation index at: https://arize-ax.mintlify.site/docs/llms.txt
> Use this file to discover all available pages before exploring further.

# Text2SQL Application for Database Querying

> Let's work through a Text2SQL use case where we are starting from scratch without a nice and clean dataset of questions, SQL queries, or expected responses.

Here are the steps:

<Steps>
  <Step title="Implement text2SQL" />

  <Step title="Setup dataset and evaluators" />

  <Step title="Run experiments" />
</Steps>

## Install dependencies

```bash theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
pip install "arize[Datasets]<8" "arize-phoenix-evals<3" openai datasets duckdb \
  pandas pyarrow
```

Set your Arize AX and OpenAI credentials as environment variables.

```bash theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
export ARIZE_SPACE_ID="<your-space-id>"
export ARIZE_API_KEY="<your-api-key>"
export ARIZE_DEVELOPER_KEY="<your-developer-key>"
export OPENAI_API_KEY="<your-openai-api-key>"
```

## Implement Text2SQL

We are going to use the NBA dataset that information from 2014 - 2018 about every game played in that span. We will use DuckDB as our database.

```python theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
import duckdb
from datasets import load_dataset

data = load_dataset("suzyanil/nba-data")["train"]

conn = duckdb.connect(database=":memory:", read_only=False)
conn.register("nba", data.to_pandas())

conn.query("SELECT * FROM nba LIMIT 5").to_df().to_dict(orient="records")[0]
```

Here's the example of one row of a game in the dataset

<Accordion title="Example data">
  `{'Unnamed: 0': 1, 'Team': 'ATL', 'Game': 1, 'Date': '10/29/14', 'Home': 'Away', 'Opponent': 'TOR', 'WINorLOSS': 'L', 'TeamPoints': 102, 'OpponentPoints': 109, 'FieldGoals': 40, 'FieldGoalsAttempted': 80, 'FieldGoals.': 0.5, 'X3PointShots': 13, 'X3PointShotsAttempted': 22, 'X3PointShots.': 0.591, 'FreeThrows': 9, 'FreeThrowsAttempted': 17, 'FreeThrows.': 0.529, 'OffRebounds': 10, 'TotalRebounds': 42, 'Assists': 26, 'Steals': 6, 'Blocks': 8, 'Turnovers': 17, 'TotalFouls': 24, 'Opp.FieldGoals': 37, 'Opp.FieldGoalsAttempted': 90, 'Opp.FieldGoals.': 0.411, 'Opp.3PointShots': 8, 'Opp.3PointShotsAttempted': 26, 'Opp.3PointShots.': 0.308, 'Opp.FreeThrows': 27, 'Opp.FreeThrowsAttempted': 33, 'Opp.FreeThrows.': 0.818, 'Opp.OffRebounds': 16, 'Opp.TotalRebounds': 48, 'Opp.Assists': 26, 'Opp.Steals': 13, 'Opp.Blocks': 9, 'Opp.Turnovers': 9, 'Opp.TotalFouls': 22}`
</Accordion>

Let's start by implementing a simple text2sql logic.

```python expandable theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
import os
import openai

client = openai.AsyncClient()

columns = conn.query("DESCRIBE nba").to_df().to_dict(orient="records")

# We will use GPT-4.1 to start
TASK_MODEL = "gpt-4.1"
CONFIG = {"model": TASK_MODEL}


system_prompt = (
    "You are a SQL expert, and you are given a single table named nba with the following columns:\n"
    f'{",".join(column["column_name"] + ": " + column["column_type"] for column in columns)}\n'
    "Write a SQL query corresponding to the user's request. Return just the query text, "
    "with no formatting (backticks, markdown, etc.)."
)


async def generate_query(input):
    response = await client.chat.completions.create(
        model=TASK_MODEL,
        temperature=0,
        messages=[
            {
                "role": "system",
                "content": system_prompt,
            },
            {
                "role": "user",
                "content": input,
            },
        ],
    )
    return response.choices[0].message.content
```

```python theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
import asyncio

query = asyncio.run(generate_query("Who won the most games?"))
print(query)
```

> SELECT Team, COUNT(\*) AS Wins&#x20;
>
> FROM nba&#x20;
>
> WHERE WINorLOSS = 'W'&#x20;
>
> GROUP BY Team&#x20;
>
> ORDER BY Wins DESC&#x20;
>
> LIMIT 1;

Let's run the query against our database now!

```python theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
def execute_query(query):
    return conn.query(query).fetchdf().to_dict(orient="records")

execute_query(query)
```

> \[\{'Team': 'GSW', 'Wins': 265}]

## Setup dataset and evaluators <a href="#evaluation" id="evaluation" />

To setup an experiment we need a dataset, task and evaluator. Let's setup each.

**Setup dataset**

```python theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
questions = [
    "Which team won the most games?",
    "Which team won the most games in 2015?",
    "Who led the league in 3 point shots?",
    "Which team had the biggest difference in records across two consecutive years?",
    "What is the average number of free throws per year?",
]
```

Let's store the data above as a versioned dataset in Arize AX.

```python theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
# Note: This example uses Python SDK v7
import json
from uuid import uuid1

import pandas as pd
from arize.experimental.datasets import ArizeDatasetsClient
from arize.experimental.datasets.utils.constants import GENERATIVE

space_id = os.environ["ARIZE_SPACE_ID"]
dataset_name = "text2sql-" + str(uuid1())[:5]

arize_client = ArizeDatasetsClient(
    developer_key=os.environ.get("ARIZE_DEVELOPER_KEY"),
    api_key=os.environ.get("ARIZE_API_KEY"),
)
# Create a dataset from a DataFrame add your own data here
test_df = pd.DataFrame([{"question": question} for question in questions])
dataset_id = arize_client.create_dataset(
    space_id=space_id,
    dataset_name=dataset_name,
    dataset_type=GENERATIVE,
    data=test_df,
)
dataset = arize_client.get_dataset(space_id=space_id, dataset_id=dataset_id)
dataset.head()
```

<Frame>
  ![](https://storage.googleapis.com/arize-phoenix-assets/assets/images/arize-docs-images/cookbooks/text2sql-1.png)
</Frame>

**Setup task**

Next, we'll define the task. The task is to generate SQL queries from natural language questions.

```python theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
async def text2sql(question):
    query = await generate_query(question)
    results = None
    error = None
    try:
        results = execute_query(query)
    except duckdb.Error as e:
        error = str(e)

    return {
        "query": query,
        "results": results,
        "error": error,
    }

# run_experiment awaits the task, so declare it async and await text2sql directly.
# Calling asyncio.run() here would fail: the runner already owns the event loop.
async def task(dataset_row):
    input = dataset_row
    return await text2sql(input["question"])
```

**Setup evaluator**

Finally, we'll define the evaluator. We'll use the following simple scoring functions to see if the generated SQL queries are correct.

```python theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
from arize.experimental.datasets.experiments.types import EvaluationResult


def _parsed(output):
    """Task output arrives as a JSON string. Returns {} if the task failed for a row."""
    if not output:
        return {}
    return json.loads(output) if isinstance(output, str) else output


# Test if there are no sql execution errors
def no_error(output):
    error = _parsed(output).get("error")
    passed = error is None
    return EvaluationResult(
        score=1.0 if passed else 0.0,
        label="no_error" if passed else "error",
        explanation="query executed cleanly" if passed else f"duckdb error: {error}",
    )


# Test if the query has results
def has_results(output):
    results = _parsed(output).get("results")
    passed = results is not None and len(results) > 0
    return EvaluationResult(
        score=1.0 if passed else 0.0,
        label="has_results" if passed else "empty",
        explanation=f"{len(results) if results else 0} row(s) returned",
    )
```

Each evaluator returns an `EvaluationResult` rather than a bare float. `label` and `explanation` are reserved columns that cannot be null, so a float-only return fails the upload.

Now let's run the experiment.

```python theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
experiment = arize_client.run_experiment(
    space_id=space_id,
    dataset_id=dataset_id,
    task=task,
    evaluators=[no_error, has_results],
    experiment_name="text2sql_first_test",
)
```

<Frame>
  ![](https://storage.googleapis.com/arize-phoenix-assets/assets/images/arize-docs-images/cookbooks/text2sql-2.png)
</Frame>

### Interpreting the results

Now that we ran the initial evaluation, it looks like three of the results are valid, one produces SQL errors, and one has no results.

The second query for `` `Which team won the most games in 2015` ``  looks for `Date LIKE '2015%'` which is not correct. The fourth query does not have TEAM in the group by clause.

Let's try to improve the prompt with few-shot examples and see if we can get better results.

```python expandable theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
samples = conn.query("SELECT * FROM nba LIMIT 1").to_df().to_dict(orient="records")[0]
sample_rows = "\n".join(
    f"{column['column_name']} | {column['column_type']} | {samples[column['column_name']]}"
    for column in columns
)
system_prompt = (
    "You are a SQL expert, and you are given a single table named nba with the following columns:\n\n"
    "Column | Type | Example\n"
    "-------|------|--------\n"
    f"{sample_rows}\n"
    "\n"
    "Write a DuckDB SQL query corresponding to the user's request. "
    "Return just the query text, with no formatting (backticks, markdown, etc.)."
)


async def generate_query(input):
    response = await client.chat.completions.create(
        model=TASK_MODEL,
        temperature=0,
        messages=[
            {
                "role": "system",
                "content": system_prompt,
            },
            {
                "role": "user",
                "content": input,
            },
        ],
    )
    return response.choices[0].message.content


print(asyncio.run(generate_query("Which team won the most games in 2015?")))
```

> SELECT Team, COUNT(\*) AS Wins FROM nba WHERE WINorLOSS = 'W' AND Date LIKE '%/15' GROUP BY Team ORDER BY Wins DESC LIMIT 1;

Looking better! Finally, let's add a scoring function that compares the results, if they exist, with the expected results. And then we can run this as another experiment and compare the results.

```python expandable theme={"theme":{"light":"github-light-default","dark":"github-dark-default"}}
# Note: This example uses Python SDK v7
from phoenix.evals.models import OpenAIModel
from phoenix.evals.classify import llm_classify


IS_SQL_EVAL_TEMPLATE = """You are a SQL expert, is the following a valid SQL query that executes without errors? Return the single workd "valid" if is valid, and "invalid" if it is not.

[BEGIN SQL QUERY]
{query}
[END SQL QUERY]
"""


def check_is_sql(output):
    query = _parsed(output).get("query")
    if not query:
        # The task produced no query for this row, so there is nothing to judge.
        # Return a result rather than None: score and label cannot be null.
        return EvaluationResult(
            score=0.0, label="invalid", explanation="task produced no query"
        )
    eval_df = llm_classify(
        dataframe=pd.DataFrame({"query": query}, index=[0]),
        template=IS_SQL_EVAL_TEMPLATE,
        model=OpenAIModel(model="gpt-4.1"),
        rails=["valid", "invalid"],
        provide_explanation=True,
    )
    label = eval_df["label"][0]
    # return score, label, explanation
    return EvaluationResult(
        score=1.0 if label == "valid" else 0.0,
        label=label,
        explanation=eval_df["explanation"][0],
    )


experiment = arize_client.run_experiment(
    space_id=space_id,
    dataset_id=dataset_id,
    task=task,
    evaluators=[no_error, has_results, check_is_sql],
    experiment_name="text2sql_test_new_prompt_and_eval-6",
)
```

<Frame>
  ![](https://storage.googleapis.com/arize-phoenix-assets/assets/images/arize-docs-images/cookbooks/text2sql-3.png)
</Frame>

You can see that the newer SQL has improved some cases, but there are still some other errors to iron out. As you experiment with different models, prompts, and techniques, you can continuously optimize your applications until they reach the performance thresholds you want.
