Building an AI Assistant with Snowflake and OpenAI: A Step-by-Step Guide

Are you looking to harness the power of AI within your data warehouse? In this tutorial, we’ll walk through creating an AI assistant that combines Snowflake’s robust data platform with OpenAI’s powerful language models. We’ll build a complete solution that includes both backend configuration in Snowflake and a user-friendly Streamlit frontend.

Prerequisites

A Snowflake account with ACCOUNTADMIN access

An OpenAI API key

Python 3.8 or later

Basic familiarity with SQL and Python

Part 1: Setting Up Snowflake Configuration

Step 1: Configure Role Permissions

First, we need to set up appropriate permissions for our warehouse. We’ll grant usage rights to both SYSADMIN and PUBLIC roles:

use role ACCOUNTADMIN;
grant usage on warehouse COMPUTE_WH to role SYSADMIN;
grant usage on warehouse COMPUTE_WH to role PUBLIC;

Step 2: Create and Configure Database

Next, we’ll create a dedicated database for our AI operations and grant necessary permissions:

use role SYSADMIN;
create database SF_LLM;

-- Grant PUBLIC role access for SELECT operations
grant usage on database SF_LLM to role PUBLIC;
grant usage on schema SF_LLM.PUBLIC to role PUBLIC;

Step 3: Configure External Access

To enable Snowflake to communicate with OpenAI’s API, we need to set up network rules and authentication:

-- Create network rule for OpenAI API access
use role SYSADMIN;
create or replace NETWORK RULE openai_network_access
mode = egress
type = host_port
value_list = ('api.openai.com');

-- Set up API key as a secret
create or replace SECRET openai_secret
  type = generic_string
  secret_string = 'API KEY';

-- Create external access integration
use role ACCOUNTADMIN;
create or replace EXTERNAL ACCESS INTEGRATION openai_access_integration
 allowed_network_rules = (openai_network_access)
 allowed_authentication_secrets = (openai_secret)
 enabled = true;

Step 4: Create the AI Function

Now we’ll create a Snowflake function that interfaces with OpenAI:

use role ACCOUNTADMIN;
create or replace function ask_mon_ami(prompt string)
returns string
language python
runtime_version=3.8
handler = 'prompt_mon_ami'
external_access_integrations=(openai_access_integration)
packages = ('openai==0.27.4','requests')
secrets = ('cred' = openai_secret )
as
$$
import _snowflake
import openai
import requests
session = requests.Session()
openai.api_key = _snowflake.get_generic_secret_string('cred')
def prompt_mon_ami(prompt):
    response = openai.ChatCompletion.create(
        model="gpt-4"
        , messages = [
        {"role": "system",
        "content": "You speak with a fake French accent"},
        { "role": "user",
        "content": prompt}
        ])
    return response.choices[0].message["content"]
$$;

Part 2: Building the Streamlit Interface

Step 1: Setting Up the Application Structure

Create a new Python file (app.py) and import the necessary libraries:

import os
import configparser
import pandas as pd
import streamlit as st
from snowflake.snowpark import Session
from snowflake.snowpark.context import get_active_session

Step 2: Implementing Snowflake Connection

Create a cached connection function to handle Snowflake authentication:

@st.cache_resource(show_spinner="Connecting to Snowflake...")
def get_session():
    try:
        return get_active_session()
    except:
        parser = configparser.ConfigParser()
        parser.read(os.path.join(os.path.expanduser('~'), ".snowsql/config"))
        section = "connections.demo_conn"
        pars = {
            "account": parser.get(section, "accountname"),
            "user": parser.get(section, "username"),
            "password": os.environ['SNOWSQL_PWD']
        }
        return Session.builder.configs(pars).create()

Step 3: Creating the AI Response Handler

Implement the function to interact with our Snowflake AI function:

@st.cache_data(show_spinner="Getting AI response...", ttl=600)
def get_ai_response(prompt):
    try:
        conn = get_session()
        query = f"SELECT ASK_MON_AMI('{prompt}')"
        result = conn.sql(query).collect()
        return result[0][0] if result else None
    except Exception as e:
        st.error(f"Error getting AI response: {str(e)}")
        return None

Step 4: Building the User Interface

Create an intuitive interface with Streamlit:

def main():
    st.title("AI Assistant Interface")
    
    chat_container = st.container()
    
    with chat_container:
        with st.form(key="prompt_form"):
            prompt = st.text_area(
                "Enter your prompt:",
                height=100,
                placeholder="Type your question here..."
            )
            submit_button = st.form_submit_button("Get Response")
            
        if submit_button and prompt:
            with st.spinner("Getting response..."):
                response = get_ai_response(prompt)
                if response:
                    with st.expander("AI Response", expanded=True):
                        st.markdown(response)
                        st.button(
                            "Copy Response",
                            key="copy_button",
                            on_click=lambda: st.write(response)
                        )
                else:
                    st.error("Failed to get response. Please try again.")

if __name__ == "__main__":
    main()

Testing the Application

To test your setup:

1. Run a simple query in Snowflake:

Beyond this example, you can leverage this infrastructure to build various other applications. Here are some interesting applications I’ve experimented with

Conclusion

We’ve successfully built a full-stack AI assistant that leverages Snowflake’s secure infrastructure and OpenAI’s powerful language models. This setup provides a foundation for building more complex AI-powered applications while maintaining data security and scalability.

Some potential enhancements you might consider:

• Adding conversation history
• Implementing role-based access control
• Adding support for different AI models
• Implementing error handling and retry mechanisms
• Adding data visualization capabilities

Remember to manage your API usage and implement appropriate rate limiting and error handling for production deployments.