Synced repo using 'sync_with_huggingface' Github Action

app.py

@@ -1,5 +1,7 @@
 from typing import Generator
-from utils import get_all_groq_model, validate_api_key, get_info, validate_uri
+from utils import validate_api_key, get_info, validate_uri, extract_code_blocks
+from langchain_community.utilities import SQLDatabase
+from var import system_prompt, markdown_info, query_output, groq_models
 import streamlit as st
 from groq import Groq
 
@@ -17,39 +19,19 @@ st.markdown("# SQL Chat")
 st.sidebar.title("Settings")
 api_key = st.sidebar.text_input("Groq API Key", type="password")
 
-models = []
-
-@st.cache_data
-def get_text_models(api_key):
-    models = get_all_groq_model(api_key=api_key)
-    vision_audio = [model for model in models if 'vision' in model or 'whisper' in model]
-    models = [model for model in models if model not in vision_audio]
-    return models
-
 # validating api_key
 if not validate_api_key(api_key):
     st.sidebar.error("Enter valid API Key")
+    model = st.sidebar.selectbox("Select Model", groq_models, disabled=True)
 else:
     st.sidebar.success("API Key is valid")
-    models = get_text_models(api_key)
-
-model = st.sidebar.selectbox("Select Model", models)
+    model = st.sidebar.selectbox("Select Model", groq_models, index=0)
 
 if st.session_state.selected_model != model:
     st.session_state.messages = []
     st.session_state.selected_model = model
 
-
 uri = st.sidebar.text_input("Enter SQL Database URI")
-db_info = {'sql_dialect': '', 'tables': '', 'tables_schema': ''}
-markdown_info = """
-**SQL Dialect**: {sql_dialect}\n
-**Tables**: {tables}\n
-**Tables Schema**:
-```sql
-{tables_schema}
-```
-"""
 
 if not validate_uri(uri):
     st.sidebar.error("Enter valid URI")
@@ -59,52 +41,20 @@ else:
     markdown_info = markdown_info.format(**db_info)
     with st.expander("SQL Database Info"):
         st.markdown(markdown_info)
+    system_prompt = system_prompt.format(markdown_info = markdown_info)
 
-system_prompt = f"""
-You are an AI assistant specialized in generating optimized SQL queries based on user instructions. \
-You have access to the database schema provided in a structured Markdown format. Use this schema to ensure \
-correctness, efficiency, and security in your SQL queries.\
-
-## SQL Database Info
-{markdown_info}
-
----
-
-## Query Generation Guidelines
-1. **Ensure Query Validity**: Use only the tables and columns defined in the schema.
-2. **Optimize Performance**: Prefer indexed columns for filtering, avoid `SELECT *` where specific columns suffice.
-3. **Security Best Practices**: Always use parameterized queries or placeholders instead of direct user inputs.
-4. **Context Awareness**: Understand the intent behind the query and generate the most relevant SQL statement.
-5. **Formatting**: Return queries in a clean, well-structured format with appropriate indentation.
-6. **Commenting**: Include comments in complex queries to explain logic when needed.
-
----
-
-## Expected Output Format
-
-The SQL query should be returned as a formatted code block:
-
-```sql
--- Get all completed orders with user details
-SELECT orders.id, users.name, users.email, orders.amount, orders.created_at
-FROM orders
-JOIN users ON orders.user_id = users.id
-WHERE orders.status = 'completed'
-ORDER BY orders.created_at DESC;
-```
-
-If the user's request is ambiguous, ask clarifying questions before generating the query.
-"""
-
-if model is not None and validate_uri(uri):
+if validate_api_key(api_key) and validate_uri(uri):
     client = Groq(
         api_key=api_key,
     )
 
+    db = SQLDatabase.from_uri(uri)
+    
+    avatar = {"user": '👨‍💻', "assistant": '🤖', "executor": '🛢'}
+
     # Display chat messages from history on app rerun
     for message in st.session_state.messages:
-        avatar = '🤖' if message["role"] == "assistant" else '👨‍💻'
-        with st.chat_message(message["role"], avatar=avatar):
+        with st.chat_message(message["role"], avatar=avatar[message["role"]]):
             st.markdown(message["content"])
 
 
@@ -118,7 +68,7 @@ if model is not None and validate_uri(uri):
     if prompt := st.chat_input("Enter your prompt here..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
 
-        with st.chat_message("user", avatar='👨‍💻'):
+        with st.chat_message("user", avatar=avatar["user"]):
             st.markdown(prompt)
 
         # Fetch response from Groq API
@@ -135,25 +85,41 @@ if model is not None and validate_uri(uri):
                         "role": m["role"],
                         "content": m["content"]
                     }
-                    for m in st.session_state.messages
+                    for m in st.session_state.messages[-8:]
                 ],
                 max_tokens=3000,
                 stream=True
             )
 
             # Use the generator function with st.write_stream
-            with st.chat_message("SQL Assistant", avatar="🤖"):
+            with st.chat_message("SQL Assistant", avatar=avatar["assistant"]):
                 chat_responses_generator = generate_chat_responses(chat_completion)
-                full_response = st.write_stream(chat_responses_generator)
+                llm_response = st.write_stream(chat_responses_generator)
+
+            with st.chat_message("SQL Executor", avatar=avatar["executor"]):
+                query = extract_code_blocks(llm_response)
+                result = db.run(query[0])
+                query_response = st.write(query_output.format(result=result))
+
         except Exception as e:
             st.error(e, icon="🚨")
 
-        # Append the full response to session_state.messages
-        if isinstance(full_response, str):
+        if len(str(result)) > 1000:
+            query_output_truncated = query_output.format(result=result)[:500]+query_output.format(result=result)[-500:]
+        else:
+            query_output_truncated = query_output.format(result=result)
+
+        # Append the llm response to session_state.messages
+        if isinstance(llm_response, str):
             st.session_state.messages.append(
-                {"role": "assistant", "content": full_response})
+                {"role": "assistant", "content": llm_response + query_output_truncated})
         else:
-            # Handle the case where full_response is not a string
-            combined_response = "\n".join(str(item) for item in full_response)
+            # Handle the case where llm_response is not a string
+            combined_response = "\n".join(str(item) for item in llm_response)
             st.session_state.messages.append(
-                {"role": "assistant", "content": combined_response})
+                {"role": "assistant", "content": combined_response + query_output_truncated})
+
+    st.sidebar.button("Clear Chat History", on_click=lambda: st.session_state.messages.clear())
+
+else:
+    st.error("Please enter valid Groq API Key and URI in the sidebar.")

utils.py

@@ -1,6 +1,7 @@
 import requests
 from langchain_community.utilities import SQLDatabase
 from langchain_community.tools.sql_database.tool import ListSQLDatabaseTool, InfoSQLDatabaseTool
+import re
 
 def get_all_groq_model(api_key:str=None) -> list:
     if api_key is None:
@@ -44,5 +45,13 @@ def get_info(uri:str) -> dict[str, str] | None:
     tables_schemas = InfoSQLDatabaseTool(db=db).invoke(access_tables)
     return {'sql_dialect': dialect, 'tables': access_tables, 'tables_schema': tables_schemas}
 
+def extract_code_blocks(text):
+    pattern = r"```(?:\w+)?\n(.*?)\n```"
+    matches = re.findall(pattern, text, re.DOTALL)
+    return matches
+
 if __name__ == "__main__":
-    print(get_all_groq_model())
+    models = (get_all_groq_model("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"))
+    vision_audio = [model for model in models if 'vision' in model or 'whisper' in model]
+    models = [model for model in models if model not in vision_audio]
+    print(models)

var.py

@@ -0,0 +1,62 @@
+groq_models = ['llama-3.3-70b-versatile', 'gemma2-9b-it', 'llama-3.2-3b-preview', 'deepseek-r1-distill-llama-70b', 'qwen-2.5-coder-32b',
+               'mixtral-8x7b-32768', 'llama-3.1-8b-instant', 'llama-3.2-1b-preview', 'allam-2-7b', 'qwen-qwq-32b', 'llama3-70b-8192',
+               'mistral-saba-24b', 'deepseek-r1-distill-qwen-32b', 'qwen-2.5-32b', 'llama-3.3-70b-specdec', 'llama3-8b-8192', 'llama-guard-3-8b']
+
+db_info = {'sql_dialect': '', 'tables': '', 'tables_schema': ''}
+
+markdown_info = """
+**SQL Dialect**: {sql_dialect}\n
+**Tables**: {tables}\n
+**Tables Schema**:
+```sql
+{tables_schema}
+```
+"""
+
+system_prompt = """
+You are an AI assistant specialized in generating optimized SQL queries based on user instructions. \
+You have access to the database schema provided in a structured Markdown format. Use this schema to ensure \
+correctness, efficiency, and security in your SQL queries.\
+
+## SQL Database Info
+{markdown_info}
+
+---
+
+## Query Generation Guidelines
+1. **Ensure Query Validity**: Use only the tables and columns defined in the schema.
+2. **Optimize Performance**: Prefer indexed columns for filtering, avoid `SELECT *` where specific columns suffice.
+3. **Security Best Practices**: Always use parameterized queries or placeholders instead of direct user inputs.
+4. **Context Awareness**: Understand the intent behind the query and generate the most relevant SQL statement.
+5. **Formatting**: Return queries in a clean, well-structured format with appropriate indentation.
+6. **Commenting**: Include comments in complex queries to explain logic when needed.
+7. **Result**: Don't return the result of the query, just the SQL query.
+8. **Optimal**: Try to generate query which is optimal and not brute force.
+9. **Single query**: Generate a best single SQL query for the user input.'
+10. **Comment**: Include comments in the query to explain the logic behind it.
+
+---
+
+## Expected Output Format
+
+The SQL query should be returned as a formatted code block:
+
+```sql
+-- Get all completed orders with user details
+-- Comment explaining the logic.
+SELECT orders.id, users.name, users.email, orders.amount, orders.created_at
+FROM orders
+JOIN users ON orders.user_id = users.id
+WHERE orders.status = 'completed'
+ORDER BY orders.created_at DESC;
+```
+
+If the user's request is ambiguous, ask clarifying questions before generating the query.
+"""
+
+query_output = """
+**The result of query execution:**
+```sql
+{result}
+```
+"""