Streamlit でSnowflake のテーブル情報のサマリを出力するアプリ

概要

Streamlit でSnowflake のテーブル情報のサマリを出力するアプリを作成する．サイドバー側で，対象のデータベース，スキーマ，テーブルを選択肢し，完了ボタンを投下すると，画面右にテーブル情報のサマリが出力される．データベースは，Snowflake Marketplace に公開されているtruestar社のJAPANESE PREFECTURE DATA事務所の統計に関するデータを用いている．

サンプルコード

import snowflake.connector
import streamlit as st
import pandas as pd

# Snowflake接続設定
conn = snowflake.connector.connect(
    user='<USERNAME>',
    password='<PASSWORD>',
    account='<ACCOUNT_NAME>',
    warehouse='<WAREHOUSE_NAME>',
    database='<DATABASE_NAME>',
    schema='<SCHEMA_NAME>'
)

def main():
    pd.set_option("display.max_colwidth",1000) 
    lang_model = "mistral-7b"
    chat(lang_model)

# Streamlitアプリケーション
def chat(lang_model):

    st.session_state.databases = {}
    st.session_state.schemas = {}
    st.session_state.tables = {}
    
    tmp = pd.read_sql(f"show databases",conn)
    databases = tmp["name"]
    st.session_state.databases = databases
    selected_database = st.sidebar.selectbox('データベース選択',st.session_state.databases)
    
    if selected_database:
        tmp = pd.read_sql(f"show schemas in database {selected_database}",conn)
        schemas = tmp["name"]
        st.session_state.schemas[selected_database] = schemas
        selected_schema = st.sidebar.selectbox('スキーマ選択',st.session_state.schemas[selected_database])
    
        if selected_schema:
            tmp = pd.read_sql(f"show objects in {selected_database}.{selected_schema}",conn)
            tables = tmp["name"]
            st.session_state.tables[f"{selected_database}.{selected_schema}"] = tables
            selected_table = st.sidebar.selectbox("テーブル選択",st.session_state.tables[f"{selected_database}.{selected_schema}"])
    
            if "messages" not in st.session_state:
                st.session_state.messages = []
            
            for message in st.session_state.messages:
                with st.chat_message(message["role"]):
                    st.markdown(message["content"])
            
            if st.sidebar.button('完了'):
                table = pd.read_sql(f"select * from {selected_database}.{selected_schema}.{selected_table} limit 100;",conn)
                data2 = ""
                columns = table.columns.values
                for i in range(0,table.shape[1]):
                    data2 = data2 + f"カラム{i}は{columns[i]}"
                    for j in range(0,table.shape[0]):
                        data2 = data2 + str(table.iloc[j,i])
                sql = "select snowflake.cortex.complete('" + lang_model + "','" + data2 + f"の情報をどんなカラムが入っているか端的に日本語で説明してください。') as response"
                sql2 = "select snowflake.cortex.count_tokens('" + lang_model + "','" + data2 + "　日本語で回答して。') * 0.12 / 1000000 as response"
                response = pd.read_sql(sql, conn)
                response2 = pd.read_sql(sql2, conn)
                
                with st.chat_message("assistant"):
                    st.markdown(response.iloc[0][0])
                    st.markdown("使用クレジット")
                    st.markdown(format(response2.iloc[0][0],'10f'))
                st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == '__main__':
    main()

こちらはアプリで表示させた画面である．