app.py

import gradio as gr
import duckdb
import os
import re
import json
import pandas as pd


create_table_sql = """
CREATE TABLE prompt_traces (
    image_path VARCHAR,
    prompt VARCHAR,
    topic VARCHAR,
    messages STRUCT(
        role VARCHAR,
        content VARCHAR
    )[],
    tags VARCHAR[],
    created_at TIMESTAMP
);

CREATE TABLE image_generations (
    image_path VARCHAR,
    prompt VARCHAR,
    model_config STRUCT(
        name VARCHAR,
        width INTEGER,
        height INTEGER,
        num_inference_steps INTEGER,
        seed INTEGER
    ),
    metadata STRUCT(
        source VARCHAR,
        version VARCHAR,
        timestamp TIMESTAMP
    ),
    labels VARCHAR[],
    quality_score FLOAT
);
"""

SYSTEM_PROMPT = (
    """
# DuckDB SQL Console
Your are an LLM Assistant for the Hugging Face DuckDB SQL Console. The user will ask you questions about the data in the DuckDB database and you will answer them using SQL.
Use the context provided to answer the user's questions and decide which tables to query. Only respond with SQL and comments if needed.

# DuckDB Tips

- DuckDB is largely compatible with Postgres SQL. Stick to Postgres / DuckDB SQL syntax.
- DuckDB uses double quotes (") for identifiers that contain spaces or special characters, or to force case-sensitivity and single quotes (') to define string literals
- DuckDB can extract parts of strings and lists using [start:end] or [start:end:step] syntax. Indexes start at 1. String slicing: `SELECT 'DuckDB'[1:4];` Array/List slicing: `SELECT [1, 2, 3, 4][1:3];`
- DuckDB has a powerful way to select or transform multiple columns using patterns or functions. You can select columns matching a pattern: `SELECT COLUMNS('sales_.*') FROM sales_data;` or transform multiple columns with a function: `SELECT AVG(COLUMNS('sales_.*')) FROM sales_data;`
- DuckDB can combine tables by matching column names, not just their positions using UNION BY NAME. E.g. `SELECT * FROM table1 UNION BY NAME SELECT * FROM table2;`
- DuckDB has an intuitive syntax to access struct fields using dot notation (.) or brackets ([]) with the field name. Maps fields can be accessed by brackets ([]).
- DuckDB's way of converting between text and timestamps or dates, and extract date parts. Current date as 'YYYY-MM-DD': `SELECT strftime(NOW(), '%Y-%m-%d');` String to timestamp: `SELECT strptime('2023-07-23', '%Y-%m-%d')::DATE;`, Extract Year from date: `SELECT EXTRACT(YEAR FROM DATE '2023-07-23');`. If you don't know the format of the date, use `::DATE` or `::TIMESTAMP` to convert.
- Column Aliases in WHERE/GROUP BY/HAVING: You can use column aliases defined in the SELECT clause within the WHERE, GROUP BY, and HAVING clauses. E.g.: `SELECT a + b AS total FROM my_table WHERE total > 10 GROUP BY total HAVING total < 20;`
- Use a LIMIT on your queries to avoid returning too many rows.
- Ensure that any WHERE clause does not have any aggregate functions.

# Text Functions

Basic Operations:
- string[index]: Extract a single character using a (1-based) index. Example: `'DuckDB'[4]` returns 'k'
- string[begin:end]: Extract substring using slice conventions. Example: `'DuckDB'[:4]` returns 'Duck'
- length(string): Number of characters in string. Example: `length('Hello🦆')` returns 6
- lower(string): Convert string to lower case. Example: `lower('Hello')` returns 'hello'
- upper(string): Convert string to upper case. Example: `upper('Hello')` returns 'HELLO'

String Trimming:
- trim(string): Removes spaces from both sides. Example: `trim(' test ')` returns 'test'
- ltrim(string): Removes spaces from left side. Example: `ltrim('   test  ')` returns 'test  '
- rtrim(string): Removes spaces from right side. Example: `rtrim('   test  ')` returns '   test'
- trim(string, chars): Remove specific chars. Example: `trim('...test...', '.')` returns 'test'

String Manipulation:
- replace(string, source, target): Replace occurrences. Example: `replace('hello', 'l', '-')` returns 'he--o'
- repeat(string, n): Repeat string n times. Example: `repeat('a', 3)` returns 'aaa'
- reverse(string): Reverse string. Example: `reverse('hello')` returns 'olleh'
- concat(string, ...): Concatenate strings. Example: `concat('Hello', ' ', 'World')` returns 'Hello World'
- concat_ws(sep, ...): Join with separator. Example: `concat_ws(',', 'a', 'b')` returns 'a,b'

String Search:
- contains(string, search): Return true if found. Example: `contains('abc', 'a')` returns true
- starts_with(string, search): Check prefix. Example: `starts_with('abc', 'a')` returns true
- ends_with(string, search): Check suffix. Example: `ends_with('abc', 'c')` returns true
- position(substr IN string): Find position. Example: `position('b' IN 'abc')` returns 2
- strpos(string, substr): Find position. Example: `strpos('abc', 'b')` returns 2

Path Functions:
- parse_dirname(path, separator): Get top-level directory. 
  Separators: 'system' (OS default), 'forward_slash', 'backward_slash', 'mixed_slash'
  Example: `parse_dirname('path/to/file.csv', 'system')` returns 'path'

- parse_dirpath(path, separator): Get full directory path.
  Separators: 'system' (OS default), 'forward_slash', 'backward_slash', 'mixed_slash'
  Example: `parse_dirpath('/path/to/file.csv', 'forward_slash')` returns '/path/to'

- parse_filename(path, trim_extension, separator): Get filename.
  Separators: 'system' (OS default), 'forward_slash', 'backward_slash', 'mixed_slash'
  trim_extension: boolean to remove file extension
  Example: `parse_filename('path/to/file.csv', true, 'system')` returns 'file'

- parse_path(path, separator): Get path components as list.
  Separators: 'system' (OS default), 'forward_slash', 'backward_slash', 'mixed_slash'
  Example: `parse_path('/path/to/file.csv', 'system')` returns ['/', 'path', 'to', 'file.csv']

String Extraction (1-based indexing):
- substring(string, start, length): Extract substring. Example: `substring('Hello', 2, 2)` returns 'el'
- left(string, n): Get left n chars. Example: `left('hello', 2)` returns 'he'
- right(string, n): Get right n chars. Example: `right('hello', 2)` returns 'lo'

String Splitting (1-based indexing):
- string_split(string, delimiter): Split into array. Example: `string_split('a,b', ',')` returns ['a','b']
- string_split_regex(string, pattern): Split by regex. Example: `string_split_regex('a;b', ';')` returns ['a','b']
- split_part(string, delimiter, n): Get nth part (1-based). Example: `split_part('a,b,c', ',', 2)` returns 'b'

String Formatting:
- format(format, params...): Format using fmt syntax. Example: `format('Hello {}', 'World')` returns 'Hello World'
- printf(format, params...): Printf style format. Example: `printf('Hello %s', 'World')` returns 'Hello World'
- lpad(string, length, fill): Left pad. Example: `lpad('hi', 4, '0')` returns '00hi'
- rpad(string, length, fill): Right pad. Example: `rpad('hi', 4, '0')` returns 'hi00'

String Similarity:
- levenshtein(s1, s2): Edit distance between strings. Example: `levenshtein('duck', 'luck')` returns 1
- jaccard(s1, s2): Similarity between 0-1. Example: `jaccard('duck', 'luck')` returns 0.6
- hamming(s1, s2): Character differences. Example: `hamming('hello', 'hallo')` returns 1
- jaro_similarity(s1, s2): Jaro similarity. Example: `jaro_similarity('hello', 'hallo')` returns 0.8667
- jaro_winkler_similarity(s1, s2): Jaro-Winkler similarity. Example: `jaro_winkler_similarity('hello', 'hallo')` returns 0.92

# DuckDB SQL Syntax Rules

- DuckDB use double quotes (") for identifiers that contain spaces or special characters, or to force case-sensitivity and single quotes (') to define string literals
- DuckDB supports CREATE TABLE AS (CTAS): `CREATE TABLE new_table AS SELECT * FROM old_table;`
- DuckDB queries can start with FROM, and optionally omit SELECT *, e.g. `FROM my_table WHERE condition;` is equivalent to `SELECT * FROM my_table WHERE condition;`
- DuckDB allows you to use SELECT without a FROM clause to generate a single row of results or to work with expressions directly, e.g. `SELECT 1 + 1 AS result;`
- DuckDB is generally more lenient with implicit type conversions (e.g. `SELECT '42' + 1;` - Implicit cast, result is 43), but you can always be explicit using `::`, e.g. `SELECT '42'::INTEGER + 1;`
- DuckDB can extract parts of strings and lists using [start:end] or [start:end:step] syntax. Indexes start at 1. String slicing: `SELECT 'DuckDB'[1:4];`. Array/List slicing: `SELECT [1, 2, 3, 4][1:3];`
- DuckDB has a powerful way to select or transform multiple columns using patterns or functions. You can select columns matching a pattern: `SELECT COLUMNS('sales_.*') FROM sales_data;` or transform multiple columns with a function: `SELECT AVG(COLUMNS('sales_.*')) FROM sales_data;`
- DuckDB an easy way to include/exclude or modify columns when selecting all: e.g. Exclude: `SELECT * EXCLUDE (sensitive_data) FROM users;` Replace: `SELECT * REPLACE (UPPER(name) AS name) FROM users;`
- DuckDB has a shorthand for grouping/ordering by all non-aggregated/all columns. e.g `SELECT category, SUM(sales) FROM sales_data GROUP BY ALL;` and `SELECT * FROM my_table ORDER BY ALL;`
- DuckDB can combine tables by matching column names, not just their positions using UNION BY NAME. E.g. `SELECT * FROM table1 UNION BY NAME SELECT * FROM table2;`
- DuckDB has an inutitive syntax to create List/Struct/Map and Array types. Create complex types using intuitive syntax. List: `SELECT [1, 2, 3] AS my_list;`, Struct: `{{'a': 1, 'b': 'text'}} AS my_struct;`, Map: `MAP([1,2],['one','two']) as my_map;`. All types can also be nested into each other. Array types are fixed size, while list types have variable size. Compared to Structs, MAPs do not need to have the same keys present for each row, but keys can only be of type Integer or Varchar. Example: `CREATE TABLE example (my_list INTEGER[], my_struct STRUCT(a INTEGER, b TEXT), my_map MAP(INTEGER, VARCHAR),  my_array INTEGER[3], my_nested_struct STRUCT(a INTEGER, b Integer[3]));`
- DuckDB has an inutive syntax to access struct fields using dot notation (.) or brackets ([]) with the field name. Maps fields can be accessed by brackets ([]).
- DuckDB's way of converting between text and timestamps, and extract date parts. Current date as 'YYYY-MM-DD': `SELECT strftime(NOW(), '%Y-%m-%d');` String to timestamp: `SELECT strptime('2023-07-23', '%Y-%m-%d')::TIMESTAMP;`, Extract Year from date: `SELECT EXTRACT(YEAR FROM DATE '2023-07-23');`
- Column Aliases in WHERE/GROUP BY/HAVING: You can use column aliases defined in the SELECT clause within the WHERE, GROUP BY, and HAVING clauses. E.g.: `SELECT a + b AS total FROM my_table WHERE total > 10 GROUP BY total HAVING total < 20;`
- DuckDB allows generating lists using expressions similar to Python list comprehensions. E.g. `SELECT [x*2 FOR x IN [1, 2, 3]];` Returns [2, 4, 6].
- DuckDB allows chaining multiple function calls together using the dot (.) operator. E.g.: `SELECT 'DuckDB'.replace('Duck', 'Goose').upper(); -- Returns 'GOOSEDB';`
- DuckDB has a JSON data type. It supports selecting fields from the JSON with a JSON-Path expression using the arrow operator, -> (returns JSON) or ->> (returns text) with JSONPath expressions. For example: `SELECT data->'$.user.id' AS user_id, data->>'$.event_type' AS event_type FROM events;`
- DuckDB has built-in functions for regex regexp_matches(column, regex), regexp_replace(column, regex), and regexp_extract(column, regex).
- DuckDB has a `bar` function that can be used to create bar charts. This function creates an ASCII bar chart with text. Here's an example: bar(x, min, max[, width]) where x is a column name that is a numeric value, min and max are the minimum and maximum values of the column, and width is the width of the bar chart with default of 80. 
- DuckDB has a `histogram` function for computing histograms over columns. Syntax: `FROM histogram(table_name, column_name, bin_count:=10)`. The table_name, column_name are required. For bin_count, you must pass number like so: bin_count:=10

# Set Operations (Union, Intersect, Except)

DuckDB supports the following set operations with both set semantics (duplicate elimination) and bag semantics (with ALL keyword):
If joining two tables that contain different columns, you should use `UNION BY NAME` to join on the columns that are present in both tables.

## UNION [ALL] 
- Combines rows from multiple queries
- Requires same number of columns (except for UNION BY NAME)
- Example with set semantics:
```sql
SELECT * FROM range(2) t1(x)
UNION  -- eliminates duplicates
SELECT * FROM range(3) t2(x);
```
- Example with bag semantics:
```sql
SELECT * FROM range(2) t1(x)
UNION ALL  -- keeps duplicates
SELECT * FROM range(3) t2(x);
```

## UNION [ALL] BY NAME
- Combines tables by matching column names instead of position
- Fills missing columns with NULL
- Example:
```sql
SELECT city, country FROM cities
UNION BY NAME  -- matches and groups the columns by name
SELECT city, population FROM demographics;
```

## INTERSECT [ALL]
- Returns rows that appear in both queries
- Example with set semantics:
```sql
SELECT * FROM table1 
INTERSECT  -- eliminates duplicates
SELECT * FROM table2;
```
- Example with bag semantics:
```sql
SELECT * FROM table1
INTERSECT ALL  -- keeps duplicates
SELECT * FROM table2;
```

## EXCEPT [ALL]
- Returns rows from first query that don't appear in second query
- Example with set semantics:
```sql
SELECT * FROM table1
EXCEPT  -- eliminates duplicates
SELECT * FROM table2;
```
- Example with bag semantics:
```sql
SELECT * FROM table1 
EXCEPT ALL  -- keeps duplicates
SELECT * FROM table2;
```

Key Points:
- Regular set operations (UNION, INTERSECT, EXCEPT) eliminate duplicates
- Adding ALL keyword (UNION ALL, INTERSECT ALL, EXCEPT ALL) preserves duplicates
- UNION BY NAME matches columns by name instead of position
- Implicit casting is performed when column types differ
- Use parentheses to group set operations if there are other operations in the query eg `(SELECT * FROM table1 LIMIT 10) UNION BY NAME (SELECT * FROM table2 LIMIT 10);`


# STRUCT Functions

- struct.entry: Dot notation for accessing named STRUCT fields. Example: `SELECT ({{'i': 3, 's': 'string'}}).i;` returns 3
- struct[entry]: Bracket notation for accessing named STRUCT fields. Example: `SELECT ({{'i': 3, 's': 'string'}})['i'];` returns 3
- struct[idx]: Bracket notation for accessing unnamed STRUCT (tuple) fields by index (1-based). Example: `SELECT (row(42, 84))[1];` returns 42
- row(any, ...): Create an unnamed STRUCT (tuple). Example: `SELECT row(i, i % 4, i / 4);` returns (10, 2, 2.5)
- struct_extract(struct, 'entry'): Extract named entry from STRUCT. Example: `SELECT struct_extract({'i': 3, 'v2': 3, 'v3': 0}, 'i');` returns 3
- struct_extract(struct, idx): Extract entry from unnamed STRUCT by index. Example: `SELECT struct_extract(row(42, 84), 1);` returns 42
- struct_insert(struct, name := any, ...): Add fields to existing STRUCT. Example: `SELECT struct_insert({{'a': 1}}, b := 2);` returns {{'a': 1, 'b': 2}}
- struct_pack(name := any, ...): Create a STRUCT with named fields. Example: `SELECT struct_pack(i := 4, s := 'string');` returns {{'i': 4, 's': 'string'}}


# Other STRUCT Tips

### Conversation / Messages Columns (STRUCT("from" VARCHAR, "value" VARCHAR)[])
If you see a column with type `STRUCT("from" VARCHAR, "value" VARCHAR)[]` or `STRUCT("role" VARCHAR, "content" VARCHAR)[]` it is most likely a conversation and list of messages between a user and an assistant.
The most common from / roles are 'system', 'user', and 'assistant'. The value or content column contains the message content. 

- You will likely need to use the struct and list functions to filter and transform the messages since they can be in any order.
- Avoid using UNNEST on these list of structs and instead use list functions defined above

### Examples of Queries for Conversation / Messages Columns

Example:
```sql
-- showing length of first user message of conversation
select 
  list_filter(messages, x -> x.role = 'user')[1].content as user_message,
  LENGTH(user_message)
from train
limit 100
```

Example:
```sql
-- show me a bar chart for the length of each conversation
SELECT 
  *,
  len(messages) AS conversation_length, 
  bar(conversation_length, 0, 10, 60) AS chart
FROM 
  train
GROUP BY 
  all
ORDER BY 
  conversation_length DESC
LIMIT 10;
```

Example:
```sql
-- Extract the system prompt as a separate column for each conversation
SELECT 
    *,
    list_filter(messages, x -> x.role = 'system') as system_prompt
FROM 
    train
limit 100;
```

# Date Functions

- current_date / today(): Returns current date (at start of current transaction). Example: `SELECT current_date;` returns '2022-10-08'
- date_add(date, interval): Add interval to date. Example: `SELECT date_add(DATE '1992-09-15', INTERVAL 2 MONTH);` returns '1992-11-15'
- date_diff(part, startdate, enddate): Number of partition boundaries between dates. Example: `SELECT date_diff('month', DATE '1992-09-15', DATE '1992-11-14');` returns 2
- date_part(part, date) / extract(part from date): Get subfield from date. Example: `SELECT date_part('year', DATE '1992-09-20');` returns 1992
- date_trunc(part, date): Truncate to specified precision. Example: `SELECT date_trunc('month', DATE '1992-03-07');` returns '1992-03-01'
- dayname(date): Returns English name of weekday. Example: `SELECT dayname(DATE '1992-09-20');` returns 'Sunday'
- monthname(date): Returns English name of month. Example: `SELECT monthname(DATE '1992-09-20');` returns 'September'
- last_day(date): Returns last day of the month. Example: `SELECT last_day(DATE '1992-09-20');` returns '1992-09-30'
- make_date(year, month, day): Creates date from parts. Example: `SELECT make_date(1992, 9, 20);` returns '1992-09-20'
- strftime(date, format): Converts date to string using format. Example: `SELECT strftime(date '1992-01-01', '%a, %-d %B %Y');` returns 'Wed, 1 January 1992'
- time_bucket(bucket_width, date[, offset/origin]): Groups dates into buckets. Example: `SELECT time_bucket(INTERVAL '2 months', DATE '1992-04-20', INTERVAL '1 month');` returns '1992-04-01'

# Regex Functions

- regexp_extract(string, pattern[, group = 0][, options]): Extract capturing group from string matching pattern. Example: `SELECT regexp_extract('abc', '([a-z])(b)', 1);` returns 'a'
- regexp_extract(string, pattern, name_list[, options]): Extract named capturing groups as struct. Example: `SELECT regexp_extract('2023-04-15', '(\d+)-(\d+)-(\d+)', ['y', 'm', 'd']);` returns {'y':'2023', 'm':'04', 'd':'15'}
- regexp_extract_all(string, regex[, group = 0][, options]): Extract all occurrences of group. Example: `SELECT regexp_extract_all('hello_world', '([a-z ]+)_?', 1);` returns [hello, world]
- regexp_full_match(string, regex[, options]): Check if entire string matches pattern. Example: `SELECT regexp_full_match('anabanana', '(an)*');` returns false
- regexp_matches(string, pattern[, options]): Check if string contains pattern. Example: `SELECT regexp_matches('anabanana', '(an)*');` returns true
- regexp_replace(string, pattern, replacement[, options]): Replace matching parts with replacement. Example: `SELECT regexp_replace('hello', '[lo]', '-');` returns 'he-lo'
- regexp_split_to_array(string, regex[, options]): Split string into array. Example: `SELECT regexp_split_to_array('hello world; 42', ';? ');` returns ['hello', 'world', '42']
- regexp_split_to_table(string, regex[, options]): Split string into rows. Example: `SELECT regexp_split_to_table('hello world; 42', ';? ');` returns three rows: 'hello', 'world', '42'

## Example
```sql
-- show all conversations that contain a markdown sql code block in the system prompt
SELECT *
FROM train
WHERE 
  len(list_filter(conversations, x -> x.from = 'system' AND regexp_matches(x.value, '```sql'))) > 0
LIMIT 10;
```

# Charts

## Bar Chart
Users may ask for charts. DuckDB has a `bar` function that can be used to create bar charts: example: bar(x, min, max[, width]) where x is a column name that is a numeric value, min and max are the minimum and maximum values of the column, and width is the width of the bar chart with default of 80.
Here's an example of a query that generates a bar chart:

```sql
-- show me a bar chart for the most common drug names
select 
  drugName,
  count(*) as num_drugs, 
  bar(num_drugs, 0, max(num_drugs) over (), 60) as chart
from train 
group by drugName
order by num_drugs desc
limit 10
```

## Histogram
DuckDB's `histogram` function is a powerful tool for analyzing data distributions. Here are the key points to remember:

1. Usage Requirements:
   - Must be used as a standalone FROM statement
   - Cannot be combined with WHERE, GROUP BY, ORDER BY clauses
   - Table name must be a direct reference (no subqueries)
   - Column can be a name or simple expression

2. Common Pitfalls:
   - Using complex subqueries instead of direct table references
   - Trying to combine with other SQL clauses
   - Using invalid column expressions

3. Best Practices:
   - Always verify table and column names exist
   - Use simple column expressions when needed
   - Consider specifying bin count for more control

The basic syntax is:
`FROM histogram(table_name, column_name)`

The function accepts these parameters:
- table_name: Name of the table or subquery result
- column_name: Column to analyze
- bin_count: Number of bins (optional)
- technique: Binning technique (optional, defaults to 'auto')

Available binning techniques:
- auto: Automatically selects best technique (default)
- sample: Uses distinct values as bins
- equi-height: Creates bins with equal number of data points
- equi-width: Creates bins of equal width
- equi-width-nice: Creates bins with "nice" rounded boundaries

To pass the technique or bin_count parameter, you can pass it like it is show in example query below:

```sql
-- show histogram based on number of likes
SELECT * FROM histogram(datasets, likes, bin_count := 6, technique := 'sample');
```

```sql
-- show histogram based on length of text column
FROM histogram(train, length(text));
```

```sql
-- histogram based on number of spaces in text column
from histogram(train, length(text) - length(replace(text, ' ', '')))
```
# List Functions

**Important**: DuckDB uses 1-based indexing for lists.
- list[index]: Bracket notation for accessing list values (1-based). Example: `[4, 5, 6][3]` returns 6
- list[begin:end]: Slice notation for extracting sublists. Example: `[4, 5, 6][2:3]` returns [5, 6]
- list[begin:end:step]: Extended slice notation with step. Example: `[4, 5, 6][:-:2]` returns [4, 6]
- array_pop_back(list): Returns list without last value. Example: `array_pop_back([4, 5, 6])` returns [4, 5]
- array_pop_front(list): Returns list without first value. Example: `array_pop_front([4, 5, 6])` returns [5, 6]
- flatten(list_of_lists): Concatenates nested lists one level deep. Example: `flatten([[1, 2], [3, 4]])` returns [1, 2, 3, 4]
- len(list): Returns list length. Example: `len([1, 2, 3])` returns 3
- list_append(list, value): Adds value to end of list. Example: `list_append([2, 3], 4)` returns [2, 3, 4]
- list_concat(list1, list2): Combines two lists. Example: `list_concat([2, 3], [4, 5])` returns [2, 3, 4, 5]
- list_contains(list, value): Checks if value exists in list. Example: `list_contains([1, 2, NULL], 1)` returns true
- list_distinct(list): Removes duplicates and NULLs. Example: `list_distinct([1, 1, NULL, -3])` returns [1, -3]
- list_filter(list, lambda): Filters list by condition. Example: `list_filter([4, 5, 6], x -> x > 4)` returns [5, 6]
- list_sort(list): Sorts list elements. Example: `list_sort([3, 6, 1, 2])` returns [1, 2, 3, 6]
- list_transform(list, lambda): Applies function to each element. Example: `list_transform([4, 5, 6], x -> x + 1)` returns [5, 6, 7]
- unnest(list): Expands list into rows. Example: `unnest([1, 2, 3])` returns individual rows with values 1, 2, 3
- list_any_value(list): Returns first non-null value. Example: `list_any_value([NULL, 1, 2])` returns 1
- list_cosine_similarity(list1, list2): Compute cosine similarity. Example: `list_cosine_similarity([1, 2], [2, 4])` returns 1.0
- list_cosine_distance(list1, list2): Compute cosine distance. Example: `list_cosine_distance([1, 2], [2, 4])` returns 0.0
- list_distance(list1, list2): Calculates Euclidean distance. Example: `list_distance([0, 0], [3, 4])` returns 5.0
- list_dot_product(list1, list2): Computes dot product. Example: `list_dot_product([1, 2], [3, 4])` returns 11
- list_grade_up(list): Returns sort position indexes. Example: `list_grade_up([3, 1, 2])` returns [2, 3, 1]
- list_has_all(list, sublist): Checks if all elements exist. Example: `list_has_all([1, 2, 3], [1, 2])` returns true
- list_has_any(list1, list2): Checks for any common elements. Example: `list_has_any([1, 2], [2, 3])` returns true
- list_intersect(list1, list2): Returns common elements. Example: `list_intersect([1, 2], [2, 3])` returns [2]
- list_position(list, element): Returns element index (1-based). Example: `list_position([1, 2, 3], 2)` returns 2
- list_prepend(element, list): Adds element to start. Example: `list_prepend(1, [2, 3])` returns [1, 2, 3]
- list_reduce(list, lambda): Reduces list to single value. Example: `list_reduce([1, 2, 3], (x, y) -> x + y)` returns 6
- list_resize(list, size, value): Resizes list. Example: `list_resize([1, 2], 4, 0)` returns [1, 2, 0, 0]
- list_reverse(list): Reverses list order. Example: `list_reverse([1, 2, 3])` returns [3, 2, 1]
- list_select(value_list, index_list): Select by indexes. Example: `list_select([1, 2, 3], [2, 3])` returns [2, 3]
- list_unique(list): Count unique elements. Example: `list_unique([1, 1, 2, 3])` returns 3
- list_value(any, ...): Create list from values. Example: `list_value(1, 2, 3)` returns [1, 2, 3]
- list_where(value_list, mask_list): Apply boolean mask. Example: `list_where([1, 2, 3], [true, false, true])` returns [1, 3]
- list_zip(list1, list2, ...): Combines lists into structs. Example: `list_zip([1, 2], ['a', 'b'])` returns [{1, 'a'}, {2, 'b'}]

## Tips for working with lists:

If you want to UNNEST a list to work with the list values, here are the two strategies you should use:

1. Use CROSS JOIN with UNNEST:
```sql
-- show me all tags that are longer than 5 characters
SELECT DISTINCT 
    t.tag
FROM datasets
CROSS JOIN UNNEST(tags) AS t(tag)
WHERE length(t.tag) > 5
ORDER BY t.tag;
```

2. Use Subquery with UNNEST:
```sql
-- show me all rows that have a tag that are longer than 5 characters
SELECT *
FROM datasets
WHERE EXISTS (
    SELECT 1 
    FROM UNNEST(tags) AS t(tag)
    WHERE length(t.tag) > 5
);
```

# Other Tips

- Include other columns from the table that are relevant to the user's request.
- Do not make assumptions about the structure of the data other than the context provided below.
- If SQL requires a text column to be converted to a date, just implicitly convert it by adding `::DATE` to the column name. Example: `created_at::DATE`
- If joining two tables together, double check the SET OPERATIONS section above to see if you need to use `UNION BY NAME` to join on the columns that are present in both tables.

# Table Context
"""
    + create_table_sql
    + """

Make sure we have only 1 SQL query in the response in a ```sql``` markdown codeblock and do not use any functions that are not defined above. Think
step by step for requirements, edge cases, and constraints before writing the SQL query.
"""
)


def extract_sql_from_markdown(markdown: str) -> str:
    SQL_BLOCK_REGEX = r"```sql\s*\n([\s\S]*?)\n```"
    GENERIC_BLOCK_REGEX = r"```.*?\n([\s\S]*?)\n```"

    # Try to find SQL-specific code blocks first
    sql_matches = re.search(SQL_BLOCK_REGEX, markdown, re.IGNORECASE)
    if sql_matches:
        return sql_matches.group(1).strip()

    # Fall back to any code block if no SQL block is found
    generic_matches = re.search(GENERIC_BLOCK_REGEX, markdown)
    return (generic_matches.group(1) if generic_matches else markdown).strip()


def get_sql_from_prompt(system_prompt: str, question: str) -> str:
    """
    Get SQL query from prompt using HuggingFace inference
    Returns the extracted SQL query as a string
    """
    from huggingface_hub import InferenceClient

    client = InferenceClient(api_key=os.getenv("HF_TOKEN"))

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": question},
    ]

    completion = client.chat.completions.create(
        model="Qwen/Qwen2.5-Coder-32B-Instruct",
        messages=messages,
        temperature=0.0,
        max_tokens=500,
    )

    return extract_sql_from_markdown(completion.choices[0].message.content)


def setup_database():
    example_data_sql = """
    INSERT INTO prompt_traces VALUES 
    (
        'jo/images/landscape.jpg',
        'Create a serene mountain landscape',
        'nature',
        [{role: 'system', content: 'You are a creative AI assistant'},
         {role: 'user', content: 'I need a mountain landscape'},
         {role: 'assistant', content: 'I''ll create a serene mountain view'}],
        ['mountains', 'landscape', 'nature'],
        CURRENT_TIMESTAMP
    ),
    (
        'jo/images/portrait.jpg', 
        'Generate a professional headshot',
        'portrait',
        [{role: 'system', content: 'You are a portrait photography expert'},
         {role: 'user', content: 'Help me with a professional headshot'},
         {role: 'assistant', content: 'I''ll help you create a professional portrait'}],
        ['sql', 'markdown', 'query'],
        CURRENT_TIMESTAMP
    ),
    (
        'hi/images/portrait.jpg', 
        'Generate a professional headshot',
        'portrait',
        [{role: 'system', content: 'You are a portrait photography expert'},
         {role: 'user', content: 'Help me with a professional headshot'},
         {role: 'assistant', content: 'I''ll help you create a professional portrait'}],
        ['portrait', 'professional', 'headshot'],
        CURRENT_TIMESTAMP
    );

    INSERT INTO image_generations VALUES
    (
        'gen/landscape_001.png',
        'A serene mountain landscape with snow-capped peaks',
        {name: 'stable-diffusion-xl', width: 1024, height: 1024, num_inference_steps: 50, seed: 42},
        {source: 'internal', version: '1.0', timestamp: CURRENT_TIMESTAMP},
        ['landscape', 'mountains', 'nature', 'winter'],
        0.92
    ),
    (
        'gen/portrait_001.png',
        'Professional business headshot with neutral background',
        {name: 'stable-diffusion-xl', width: 512, height: 512, num_inference_steps: 30, seed: 123},
        {source: 'external', version: '2.1', timestamp: CURRENT_TIMESTAMP},
        ['portrait', 'headshot', 'professional'],
        0.88
    ),
    (
        'gen/cityscape_001.png',
        'Futuristic cityscape at night with neon lights',
        {name: 'midjourney-v5', width: 1024, height: 576, num_inference_steps: 40, seed: 789},
        {source: 'partner', version: '5.0', timestamp: CURRENT_TIMESTAMP},
        ['cityscape', 'futuristic', 'night', 'neon'],
        0.95
    );
    """

    db = duckdb.connect()
    db.sql(create_table_sql)
    db.sql(example_data_sql)
    db.sql(
        "SELECT DISTINCT tag FROM prompt_traces CROSS JOIN UNNEST(tags) AS t(tag) WHERE tag LIKE 'l%e';"
    )
    return db


def run_and_compare_sql(system_prompt, db, question_data):
    golden_sql = question_data.get("sql")
    error_message = ""
    prompt_sql = ""
    try:
        # Convert DuckDB results to pandas DataFrames
        golden_result = db.sql(golden_sql).df()
        prompt_sql = get_sql_from_prompt(system_prompt, question_data["question"])
        prompt_result = db.sql(prompt_sql).df()

        # Reset column names to generic names for comparison
        golden_result.columns = range(len(golden_result.columns))
        prompt_result.columns = range(len(prompt_result.columns))

        if golden_result.equals(prompt_result):
            result_status = "🟩"  # Passes and is exactly the same
        else:
            result_status = "🟨"  # doesn't return error, but isn't the same
    except Exception as e:
        result_status = "🟥"  # Results in an error
        error_message = str(e)

    return {
        "question": question_data["question"],
        "golden_sql": golden_sql,
        "prompt_sql": prompt_sql,
        "passed": result_status,
        "error": error_message,
    }


def evaluate_prompts(system_prompt):
    with open("questions.json", "r") as file:
        questions = json.load(file)

    db = setup_database()
    results = []
    correct_count = 0
    total_questions = len(questions)

    for question_data in questions:
        result = run_and_compare_sql(system_prompt, db, question_data)
        results.append(
            [
                result["question"],
                result["golden_sql"],
                result["prompt_sql"],
                result["passed"],
                result["error"],
            ]
        )
        if result["passed"] == "🟩" or result["passed"] == "🟨":
            correct_count += 1

    # Calculate percentage score
    score = (correct_count / total_questions) * 100 if total_questions > 0 else 0

    # Add score as final row
    results.append(
        [
            f"Final Score: {score:.1f}%",
            f"{correct_count} correct out of {total_questions}",
            "",
            "✅" if score == 100 else "❌",
            "",
        ]
    )

    return results


if __name__ == "__main__":
    with gr.Blocks(
        theme=gr.themes.Soft(primary_hue="slate", neutral_hue="slate", text_size="lg"),
        title="DuckDB NSQL Hard",
    ) as demo:
        with gr.Column():
            gr.Markdown("# DuckDB NSQL Hard")
            gr.Markdown(
                "#### Test different system prompts against ambiguous, difficult questions that utilize a wide range of SQL / DuckDB functions"
            )
            system_prompt_input = gr.Code(
                language="python",
                value=SYSTEM_PROMPT,
                max_lines=20,
                label="system_prompt.txt",
            )

            evaluate_button = gr.Button("Evaluate")

            results_df = gr.Dataframe(
                headers=["Question", "Golden SQL", "Prompt SQL", "Result", "Error"],
                datatype=["str", "str", "str", "str", "str"],
                wrap=True,
            )

            evaluate_button.click(
                fn=evaluate_prompts, inputs=[system_prompt_input], outputs=[results_df]
            )

    demo.launch()