Update app.py

app.py

@@ -1,149 +1,56 @@
+# Questions:
+# ==========
+# Show me the total number of entries in the first table
+# Select top 10 customers from Canada with highest sum of C_ACCTBAL value, in descending order
+# Show me the total of Customers per Nation, in ascending order
+# Show me a query that lists totals for extended price, discounted extended price, discounted extended price plus tax, average quantity, average extended price, and average discount. These aggregates are grouped by RETURNFLAG and LINESTATUS, and listed in ascending order of RETURNFLAG and LINESTATUS. A count of the number of line items in each group is included
+
 import os
-import snowflake.connector
-import replicate
-import re
-import pandas as pd
 import streamlit as st
-
-# Snowflake connection parameters
-ACCOUNT = "EZ97576.ap-southeast-1"
-USER = "sureshsnowflake"
-PASSWORD = "Slavia@123"
-WAREHOUSE = "COMPUTE_WH"
-DATABASE = "SNOWFLAKE_SAMPLE_DATA"
-SCHEMA = "TPCDS_SF100TCL"
-
-# Replicate API key
-os.environ['REPLICATE_API_TOKEN'] = 'r8_E7Rn49bbi2O33bztSMYLKyqvWmo68mZ1Tg8M0'
-
-
-def interact_with_replicate(prompt):
-    response = ""
-    for event in replicate.stream(
-        "snowflake/snowflake-arctic-instruct",
-        input={
-            "prompt": prompt,
-            "max_new_tokens": 250
-        },
-    ):
-        response += str(event)
-
-    # Use regular expressions to extract SQL statements
-    sql_statements = re.findall(
-        r"(SELECT.*?;|INSERT.*?;|UPDATE.*?;|DELETE.*?;|CREATE.*?;|ALTER.*?;|DROP.*?;)", response, re.DOTALL | re.IGNORECASE)
-
-    # Join the SQL statements into a single string
-    return "\n".join(sql_statements)
-
-
-def get_snowflake_connection():
-    return snowflake.connector.connect(
-        user=USER,
-        password=PASSWORD,
-        account=ACCOUNT,
-        warehouse=WAREHOUSE,
-        database=DATABASE,
-        schema=SCHEMA
-    )
-
-
-def fetch_ddl_for_all_tables():
-    conn = get_snowflake_connection()
-    # Add headers
-    ddl_data = ["table_name, column_name, data_type, is_nullable\n"]
-
-    try:
-        cur = conn.cursor()
-        table_names_query = f"""
-        SELECT TABLE_NAME
-        FROM {DATABASE}.INFORMATION_SCHEMA.TABLES
-        WHERE TABLE_SCHEMA = '{SCHEMA}'
-        """
-        cur.execute(table_names_query)
-        table_names = cur.fetchall()
-
-        for table_name in table_names:
-            table_name = table_name[0]
-            ddl_query = f"""
-            SELECT table_name, column_name, data_type, is_nullable 
-            FROM {DATABASE}.INFORMATION_SCHEMA.COLUMNS 
-            WHERE table_name = '{table_name}' AND table_schema = '{SCHEMA}'
-            """
-            cur.execute(ddl_query)
-            ddl_result = cur.fetchall()
-            if ddl_result:
-                ddl = "\n".join(
-                    [f"{row[0]}, {row[1]}, {row[2]}, {row[3]}" for row in ddl_result])
-            else:
-                ddl = f"No DDL found for table {table_name}"
-            ddl_data.append(f"-- DDL for table {table_name} --\n{ddl}\n\n")
-
-        with open('sample_ddl.txt', 'w') as file:
-            file.writelines(ddl_data)
-
-        st.success("DDLs written to sample_ddl.txt")
-    finally:
-        cur.close()
-        conn.close()
-
-
-def generate_sql_query(sample_message):
-    with open('sample_ddl.txt', 'r') as file:
-        ddl_commands = file.read()
-
-    instruction = "Read the all provided ddl statements and work on the statement. if any logical questions asked,find relattion between tables based on primarykey and foreign key difinition use snowflake supported functions to provide snowflake sql construct. validate functions used on coulmn data types,amiguity in joins before present also add as many description columns as possible. Respond only with the SQL query without any explanations or contextual details. if the ask is to find or provide then consider as a request for writing sql construct"
-    combined_input = ddl_commands + sample_message + " " + instruction
-
-    response = interact_with_replicate(combined_input)
-
-    with open('generated_query.sql', 'w') as file:
-        file.write(response)
-
-    st.success("Response from Replicate has been written to 'generated_query.sql'")
-
-
-def execute_generated_sql():
-    with open('generated_query.sql', 'r') as file:
-        generated_sql = file.read()
-
-    # Print the generated SQL for inspection
-    st.text_area("Generated SQL query", generated_sql, height=200)
-
-    conn = get_snowflake_connection()
-
-    try:
-        cur = conn.cursor()
-        cur.execute(generated_sql)
-
-        if cur.description is not None:
-            result = cur.fetchall()
-            columns = [desc[0] for desc in cur.description]
-            df = pd.DataFrame(result, columns=columns)
-            st.write("Result from executed SQL query:")
-            st.dataframe(df)
-        else:
-            st.write(
-                "The executed SQL did not return any results or is not a SELECT query.")
-    finally:
-        cur.close()
-        conn.close()
-
-
-def main():
-    st.title("Snowflake and Replicate Integration")
-
-    st.header("Generate SQL Query and Execute")
-    sample_message = st.text_area(
-        "Enter your message for generating SQL query", height=100)
-    if st.button("Generate SQL"):
-        generate_sql_query(sample_message)
-        st.success("SQL Query generated successfully. You can now execute it.")
-
-    if st.button("Execute SQL"):
-        execute_generated_sql()
-
-
-if __name__ == "__main__":
-    # Fetch DDLs for all tables automatically before starting the app
-    fetch_ddl_for_all_tables()
-    main()
+import pandas
+from snowflake.snowpark import Session
+from langchain_community.utilities import SQLDatabase
+from langchain.chains import create_sql_query_chain
+from transformers import AutoTokenizer
+from langchain_community.llms import Replicate
+from langchain_core.prompts import PromptTemplate
+
+st.set_page_config(page_title="Snowflake Arctic", page_icon="🤖")
+@st.cache_resource(show_spinner="Connecting...")
+def getSession():
+    section = st.secrets[f"connections_snowflake"]
+    pars = {
+        "account": section["account"],
+        "user": section["user"],
+        "password": section["password"],
+        "database": section["database"],
+        "schema": section["schema"],
+        "warehouse": section["warehouse"],
+        "role": section["role"]
+    }
+    session = Session.builder.configs(pars).create()
+
+    url = (f"snowflake://{pars['user']}:{pars['password']}@{pars['account']}"
+        + f"/{pars['database']}/{pars['schema']}"
+        + f"?warehouse={pars['warehouse']}&role={pars['role']}")
+    db = SQLDatabase.from_uri(url)
+    
+    os.environ['REPLICATE_API_TOKEN'] = st.secrets["REPLICATE_API_TOKEN"]
+    llm = Replicate(model="snowflake/snowflake-arctic-instruct", model_kwargs={"temperature": 0.75, "top_p": 1},)
+    chain = create_sql_query_chain(llm, db)
+    return session, db, chain
+
+st.title("❄️ Snowflake Arctic with Replicate")
+st.write("Returns and runs queries from questions in natural language.")
+
+session, db, chain = getSession()
+
+#user_query = st.chat_input("Type your message here...")
+
+user_query = st.sidebar.text_area("Ask a question:", value="Show me the total number of entries in the first table")
+sql = chain.invoke({"question": user_query}).rstrip(';')
+
+tabQuery, tabData, tabLog = st.tabs("Query", "Data", "Log")
+tabQuery.code(sql, language="sql")
+tabData.dataframe(session.sql(sql))
+tabLog.code(db.table_info, language="sql")