Upload app.py

app.py

@@ -0,0 +1,101 @@
+from dotenv import find_dotenv, load_dotenv
+load_dotenv(find_dotenv())
+import os
+import pandas as pd
+import streamlit as st
+from langchain_core.prompts import PromptTemplate
+from pandasql import sqldf
+from groq import Groq
+
+
+template = """You are a powerful text-to-SQL model. Your job is to answer questions about a database. You are given a question and context regarding one or more tables. Dont add \n characters.
+Do not include "SELECT short\_name, long\_name" this type of queries which have backslash in them.
+You must output the SQL query that answers the question in a single line.
+
+### Input:
+`{question}`
+
+### Context:
+`{context}`
+
+### Response:
+"""
+
+
+
+prompt = PromptTemplate.from_template(template=template)
+
+client = Groq(
+    api_key=os.getenv("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"),
+)
+
+def groq_infer(prompt):
+    chat_completion = client.chat.completions.create(
+    messages=[
+        {
+            "role": "user",
+            "content": prompt,
+        }
+    ],
+    model="mixtral-8x7b-32768",
+)
+    print(chat_completion.choices[0].message.content)
+    return chat_completion.choices[0].message.content
+
+# 1. Create cache_resource - To load the model
+# infer - pipeline -> pipe()
+def main():
+    st.set_page_config(page_title="Ask anything about database", page_icon="📊", layout="wide")
+    st.title("SQL Engineer")
+
+    col1, col2 = st.columns([2, 3])
+
+    with col1:
+        uploaded_file = st.file_uploader("Upload a CSV file", type="csv")
+        if uploaded_file is not None:
+            df = pd.read_csv(uploaded_file, encoding="latin1")
+            df.columns = df.columns.str.replace(r"[^a-zA-Z0-9_]", "", regex=True)
+            st.write("Here's a preview of your uploaded file:")
+            st.dataframe(df)
+
+            context = pd.io.sql.get_schema(df.reset_index(), "df").replace('"', "")
+            st.write("SQL Schema:")
+            st.code(context)
+
+    with col2:
+        if uploaded_file is not None:
+            question = st.text_input("Write a question about the data", key="question")
+
+            if st.button("Get Answer", key="get_answer"):
+                if question:
+                    attempt = 0
+                    max_attempts = 5
+                    while attempt < max_attempts:
+                        try:
+                            input = {"context": context, "question": question}
+                            formatted_prompt = prompt.invoke(input=input).text
+                            response = groq_infer(formatted_prompt)
+                            final = response.replace("`", "").replace("sql", "").strip()
+                            st.text("Query performed")
+                            st.code(final)
+                            result = sqldf(final, locals())
+                            st.write("Answer:")
+                            st.dataframe(result)
+                            break
+                        except Exception as e:
+                            attempt += 1
+                            st.error(
+                                f"Attempt {attempt}/{max_attempts} failed. Retrying..."
+                            )
+                            if attempt == max_attempts:
+                                st.error(
+                                    "Unable to get the correct query, refresh app or try again later."
+                                )
+                            continue
+
+                else:
+                    st.warning("Please enter a question before clicking 'Get Answer'.")
+
+
+if __name__ == "__main__":
+    main()