utils/sql_utils.py

import re
import streamlit as st
import pandas as pd

def extract_sql_command(text):
    """
    Extracts the SQL command enclosed within ```sql ``` delimiters from a given string.

    Args:
        text: The input string containing the SQL command.

    Returns:
        The extracted SQL command as a string, or None if no SQL command is found.
    """
    pattern = r"```sql\s*([\s\S]*?)\s*```"
    match = re.search(pattern, text)
    if match:
        return match.group(1).strip()
    else:
        return None

def create_schema(dataframes):
    schema = ""
    for df_name, df in dataframes.items():
        schema += f"### {df_name}\n"
        schema += df.head(3).to_markdown(index=False)
        schema += "\n\nRows: " + str(df.shape[0]) + ", Columns: " + str(df.shape[1]) + "\n\n---\n\n"
    return schema


@st.cache_resource
def load_defaultdb_schema_text():
    with open("static/database_scema.txt", "r", encoding="utf-8") as file:
        return file.read()


@st.cache_resource
def load_defaultdb_queries():
    with open("static/default_questions.txt", "r", encoding="utf-8") as file:
        return file.read()


@st.cache_data
def convert_df(df):
    # IMPORTANT: Cache the conversion to prevent computation on every rerun
    return df.to_csv().encode("utf-8")


# Load CSV files into pandas default_dfs
@st.cache_resource
def load_data():
    # text-to-sql-streamlit\static\df_Customers.csv
    df_customers = pd.read_csv("static/df_Customers.csv")
    df_order_items = pd.read_csv("static/df_OrderItems.csv")
    df_orders = pd.read_csv("static/df_Orders.csv")
    df_payments = pd.read_csv("static/df_Payments.csv")
    df_products = pd.read_csv("static/df_Products.csv")
    return {
        "customers": df_customers,
        "order_items": df_order_items,
        "orders": df_orders,
        "payments": df_payments,
        "products": df_products,
    }