Create app.py

app.py

@@ -0,0 +1,149 @@
+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()