Text-to-SQL LLM App with Snowflake Cortex

Snowflake’s Cortex Analyst is an AI feature in Snowflake that allows you quickly create apps that translate natural language into SQL queries.

Not a medium member? Click here to read for free!

In this article, I’ll replicate a text-to-SQL app I built a few weeks ago for a small hackathon, although I have made minor changes in my approach here such as using the semantic model generator. It will have the following:

• Public data sourced from the City of Toronto: Short Term Rentals Registration
• Snowflake table to store this data
• Streamlit user interface, deployed within Snowflake, with which to query Cortex Analyst API
• A semantic model (or data dictionary) to guide Cortex Analyst

I’m assuming you have familiarity with the Snowflake interface so there will not be detailed breakdowns of basic steps. This demo is adapted from Snowflake Quickstart documentation (refer to references at the bottom for more).

Set Up Cortex

Ideally, you should create a cortex user role for your deployment:

USE ROLE ACCOUNTADMIN;

CREATE ROLE cortex_user_role;
GRANT DATABASE ROLE SNOWFLAKE.CORTEX_USER TO ROLE cortex_user_role;

GRANT ROLE cortex_user_role TO USER some_user;

In my case I have an existing database called supernova and I have created a schema called housing.

USE ROLE ACCOUNTADMIN;
-- Grant permissions to cortex user role
GRANT USAGE ON DATABASE SUPERNOVA TO ROLE CORTEX_USER_ROLE;
GRANT USAGE ON SCHEMA SUPERNOVA.HOUSING TO ROLE CORTEX_USER_ROLE;
GRANT CREATE STREAMLIT ON SCHEMA SUPERNOVA.HOUSING TO ROLE CORTEX_USER_ROLE;
GRANT CREATE STAGE ON SCHEMA SUPERNOVA.HOUSING TO ROLE CORTEX_USER_ROLE;
GRANT USAGE ON WAREHOUSE COMPUTE_WH TO ROLE CORTEX_USER_ROLE;
--use new role
USE ROLE CORTEX_USER_ROLE;

Take note that Cortex is not natively available in every Snowflake region. Check this snowflake documentation for the updated list of regions or use Cross Region Inference to enable Cortex in your region. My Snowflake region doesn’t natively support Cortex so I had to set up cross region inference:

USE ROLE ACCOUNTADMIN;
ALTER ACCOUNT SET CORTEX_ENABLED_CROSS_REGION = 'AWS_US';

Set Up Data in Snowflake

Download the .xlsx file and insert into a table of your choice in your Snowflake database. I have created a database called supernovaand a schema called housing. From the Snowflake UI I’ll create a new table in the housing schema from the file downloaded.

Create table from file

Data preview

Set Up Semantic Model

I decided to use the Semantic Model Generator to create the semantic model for this app. I set up the generator locally and started the streamlit app.

Semantic Model Generator home screen

After selecting create a new semantic model and filling in the database details, it generated a draft model for me to review and validate. I made minor additions to the draft and clicked validate.

Then I used the Cortex chat interface on the right side of the generator app to test the draft semantic model and saved any queries I approved as verified queries in the semantic model.

Review semantic model, test and save verified queries

Set Up Chat App with Streamlit in Snowflake

Create a new streamlit app

Switch to your cortex user role and create a new streamlit app.

Create Streamlit app

In the Snowflake UI, navigate to the stage auto-generated when you created the Streamlit app and upload your semantic model file. The stage may have a long alphanumeric name.

If you’re unable to see the file upload button, click Enable Directory Table so you can see the list of files in the stage.

Upload semantic model to streamlit stage

Files in the stage

So the way this stage has been auto-named is a bit problematic when I want to run code like this:

list @SUPERNOVA.HOUSING."NXNI4G2I0WAR7JCE (Stage)"

Syntax error: unexpected '('. (line 31)

So I’m inclined to editing the stage name:

alter stage SUPERNOVA.HOUSING."NXNI4G2I0WAR7JCE (Stage)" rename to SUPERNOVA.HOUSING.APP_STAGE;

Then I grab the streamlit app name by taking it from the output of this query:

show streamlits in housing;

streamlit apps in housing schema

Then I modify the stage for the streamlit app to the new stage name:

alter streamlit NXNI4G2I0WAR7JCE set root_location = @SUPERNOVA.HOUSING.APP_STAGE;

--confirm changes
describe streamlit NXNI4G2I0WAR7JCE;

Streamlit app details showing stage name change

Now, from Projects > Streamlit go back to your Streamlit app, click Edit and paste the code below:

import _snowflake
import json
import streamlit as st
import time
from snowflake.snowpark.context import get_active_session

DATABASE = "SUPERNOVA"
SCHEMA = "HOUSING"
STAGE = "APP_STAGE"
FILE = "semantic_model.yaml"
session = get_active_session()

def translate():
    supported_languages = {'German':'de','French':'fr','Korean':'ko','Portuguese':'pt','English':'en','Italian':'it','Russian':'ru','Swedish':'sv','Spanish':'es','Japanese':'ja','Polish':'pl'}
    with st.container():
        col1, col2 = st.columns(2)
        with col1:
            st.title("Translator")
        with col2:
            # Image credit Streamlit
            st.image(image="streamlit-bot-icon.png",width=70)
        #st.header("Translate With Snowflake Cortex")
        col3,col4 = st.columns(2)
        with col3:
            from_language = st.selectbox('From',dict(sorted(supported_languages.items())))
        with col4:
            to_language = st.selectbox('To',dict(sorted(supported_languages.items())))
        entered_text = st.text_area("Enter text",label_visibility="hidden",height=300,placeholder='Enter text. For example, a promo call to action.')
        btn_translate = st.button("Translate",type="primary")
        if entered_text and btn_translate:
          entered_text = entered_text.replace("'", "\\'")
          cortex_response = session.sql(f"select snowflake.cortex.translate('{entered_text}','{supported_languages[from_language]}','{supported_languages[to_language]}') as response").to_pandas().iloc[0]['RESPONSE']
          st.write(cortex_response)

def send_message(prompt: str) -> dict:
    """Calls the REST API and returns the response."""
    request_body = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt
                    }
                ]
            }
        ],
        "semantic_model_file": f"@{DATABASE}.{SCHEMA}.{STAGE}/{FILE}",
    }
    resp = _snowflake.send_snow_api_request(
        "POST",
        f"/api/v2/cortex/analyst/message",
        {},
        {},
        request_body,
        {},
        30000,
    )
    if resp["status"] < 400:
        return json.loads(resp["content"])
    else:
        raise Exception(
            f"Failed request with status {resp['status']}: {resp}"
        )

def process_message(prompt: str) -> None:
    """Processes a message and adds the response to the chat."""
    st.session_state.messages.append(
        {"role": "user", "content": [{"type": "text", "text": prompt}]}
    )
    with st.chat_message("user"):
        st.markdown(prompt)
    with st.chat_message("assistant"):
        with st.spinner("Generating response..."):
            response = send_message(prompt=prompt)
            content = response["message"]["content"]
            display_content(content=content)
    st.session_state.messages.append({"role": "assistant", "content": content})

def display_content(content: list, message_index: int = None) -> None:
    """Displays a content item for a message."""
    message_index = message_index or len(st.session_state.messages)
    for item in content:
        if item["type"] == "text":
            st.markdown(item["text"])
        elif item["type"] == "suggestions":
            with st.expander("Suggestions", expanded=True):
                for suggestion_index, suggestion in enumerate(item["suggestions"]):
                    if st.button(suggestion, key=f"{message_index}_{suggestion_index}"):
                        st.session_state.active_suggestion = suggestion
        elif item["type"] == "sql":
            with st.expander("SQL Query", expanded=False):
                st.code(item["statement"], language="sql")
            with st.expander("Results", expanded=True):
                with st.spinner("Running SQL..."):
                    session = get_active_session()
                    df = session.sql(item["statement"]).to_pandas()
                    if len(df.index) > 1:
                        data_tab, line_tab, bar_tab = st.tabs(
                            ["Data", "Line Chart", "Bar Chart"]
                        )
                        data_tab.dataframe(df)
                        if len(df.columns) > 1:
                            df = df.set_index(df.columns[0])
                        with line_tab:
                            st.line_chart(df)
                        with bar_tab:
                            st.bar_chart(df)
                    else:
                        st.dataframe(df)

def query_builder():
    col1, col2 = st.columns(2)
    with col1:
        st.title("Cortex Bot")
    with col2:
        # Image credit Streamlit
        st.image(image="streamlit-bot-icon.png",width=70)

    st.markdown(f"Semantic Model: `{FILE}`")

    if "messages" not in st.session_state:
        st.session_state.messages = []
        st.session_state.suggestions = []
        st.session_state.active_suggestion = None

    for message_index, message in enumerate(st.session_state.messages):
        with st.chat_message(message["role"]):
            display_content(content=message["content"], message_index=message_index)

    if user_input := st.chat_input("What is your question?"):
        process_message(prompt=user_input)

    if st.session_state.active_suggestion:
        process_message(prompt=st.session_state.active_suggestion)
        st.session_state.active_suggestion = None

page_names_to_funcs = {
    "Conversational Analytics": query_builder,
    "Translate": translate,
}

selected_page = st.sidebar.selectbox("Select", page_names_to_funcs.keys())
page_names_to_funcs[selected_page]()

In the send_message function, notice the API call to Snowflake Cortex which also supplies the semantic model file for response synthesis.

User Interface

Toggle the tables in the results container to review the auto generated charts:

Chart generated from the query by Cortex

Ask it a couple more questions and evaluate the performance. I think it’s doing really well and it performs even better with a richly defined semantic model.

Benefits

I see how a well developed text-to-SQL chat assistant will be beneficial to users:

• Non-technical users can ask questions in regular English, just as they would to a data analyst
• Certain dashboard projects will not need to be started because a semantic model can be developed over the required tables and used to enhance a Cortex Analyst app. For dashboards that may only be used infrequently, this reduces manpower requirements where there are competing priorities.
• For organizations that already use Snowflake, building a Streamlit app within your Snowflake organization simply leverages the capabilities already available to you and reduces the need to built separate systems that require new infrastructure, data movements or permissions.

Bonus

I left a small piece of code to demonstrate how translation with Cortex works. On the Streamlit UI go to the sidebar and select the translate option and try to translate anything.

References

1. Cortex Analyst | Snowflake Documentation
2. Creating Semantic Models for Cortex Analyst
3. Cross-region inference | Snowflake Documentation
4. A Getting Started Guide With Snowflake Arctic and Snowflake Cortex
5. Snowflake AI and ML | Snowflake Documentation
6. GitHub — Snowflake-Labs/semantic-model-generator