Synced repo using 'sync_with_huggingface' Github Action

app.py

@@ -1,5 +1,5 @@
 from typing import Generator
-from utils import validate_api_key, get_info, validate_uri, extract_code_blocks
+from utils import validate_api_key, get_info, validate_uri, extract_code_blocks, get_info_sqlalchemy
 from langchain_community.utilities import SQLDatabase
 from var import system_prompt, markdown_info, query_output, groq_models
 import streamlit as st
@@ -10,6 +10,7 @@ st.set_page_config(layout="wide")
 # Initialize chat history and selected model
 if "messages" not in st.session_state:
     st.session_state.messages = []
+    st.session_state.sql_result = []
 
 if "selected_model" not in st.session_state:
     st.session_state.selected_model = None
@@ -29,6 +30,7 @@ else:
 
 if st.session_state.selected_model != model:
     st.session_state.messages = []
+    st.session_state.sql_result = []
     st.session_state.selected_model = model
 
 uri = st.sidebar.text_input("Enter SQL Database URI")
@@ -37,7 +39,7 @@ if not validate_uri(uri):
     st.sidebar.error("Enter valid URI")
 else:
     st.sidebar.success("URI is valid")
-    db_info = get_info(uri)
+    db_info = get_info_sqlalchemy(uri)
     markdown_info = markdown_info.format(**db_info)
     with st.expander("SQL Database Info"):
         st.markdown(markdown_info)
@@ -53,9 +55,12 @@ if validate_api_key(api_key) and validate_uri(uri):
     avatar = {"user": '👨‍💻', "assistant": '🤖', "executor": '🛢'}
 
     # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
+    for i, message in enumerate(st.session_state.messages):
         with st.chat_message(message["role"], avatar=avatar[message["role"]]):
             st.markdown(message["content"])
+        if (i+1)%2 == 0:
+            with st.chat_message("SQL Executor", avatar=avatar["executor"]):
+                st.markdown(st.session_state.sql_result[i//2])
 
 
     def generate_chat_responses(chat_completion) -> Generator[str, None, None]:
@@ -109,17 +114,19 @@ if validate_api_key(api_key) and validate_uri(uri):
         else:
             query_output_truncated = query_output.format(result=result)
 
+        st.session_state.sql_result.append(query_output_truncated)
+
         # Append the llm response to session_state.messages
         if isinstance(llm_response, str):
             st.session_state.messages.append(
-                {"role": "assistant", "content": llm_response + query_output_truncated})
+                {"role": "assistant", "content": llm_response})
         else:
             # 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 + query_output_truncated})
+                {"role": "assistant", "content": combined_response})
 
-    st.sidebar.button("Clear Chat History", on_click=lambda: st.session_state.messages.clear())
+    st.sidebar.button("Clear Chat History", on_click=lambda: st.session_state.messages.clear() and st.session_state.sql_result.clear())
 
 else:
     st.error("Please enter valid Groq API Key and URI in the sidebar.")

utils.py

@@ -1,9 +1,21 @@
 import requests
 from langchain_community.utilities import SQLDatabase
 from langchain_community.tools.sql_database.tool import ListSQLDatabaseTool, InfoSQLDatabaseTool
+from sqlalchemy import (
+    create_engine,
+    MetaData,
+    inspect,
+    Table,
+    select,
+    distinct
+)
+from sqlalchemy.schema import CreateTable
+from sqlalchemy.exc import ProgrammingError
+from sqlalchemy.engine import Engine
 import re
 
 def get_all_groq_model(api_key:str=None) -> list:
+    """Uses Groq API to fetch all the available models."""
     if api_key is None:
         raise ValueError("API key is required")
     url = "https://api.groq.com/openai/v1/models"
@@ -21,6 +33,7 @@ def get_all_groq_model(api_key:str=None) -> list:
     return model_ids
 
 def validate_api_key(api_key:str) -> bool:
+    """Validates the Groq API key using the get_all_groq_model function."""
     if len(api_key) == 0:
         return False
     try:
@@ -30,6 +43,7 @@ def validate_api_key(api_key:str) -> bool:
         return False
 
 def validate_uri(uri:str) -> bool:
+    """Validates the SQL Database URI using the SQLDatabase.from_uri function."""
     try:
         SQLDatabase.from_uri(uri)
         return True
@@ -37,6 +51,7 @@ def validate_uri(uri:str) -> bool:
         return False
 
 def get_info(uri:str) -> dict[str, str] | None:
+    """Gets the dialect name, accessible tables and table schemas using the SQLDatabase toolkit"""
     db = SQLDatabase.from_uri(uri)
     dialect = db.dialect
     # List all the tables accessible to the user.
@@ -45,10 +60,95 @@ 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 get_sample_rows(engine:Engine, table:Table, row_count: int = 3) -> str:
+    """Gets the sample rows of a table using the SQLAlchemy engine"""
+    # build the select command
+    command = select(table).limit(row_count)
+
+    # save the columns in string format
+    columns_str = "\t".join([col.name for col in table.columns])
+
+    try:
+        # get the sample rows
+        with engine.connect() as connection:
+            sample_rows_result = connection.execute(command)  # type: ignore
+            # shorten values in the sample rows
+            sample_rows = list(
+                map(lambda ls: [str(i)[:100] for i in ls], sample_rows_result)
+            )
+
+        # save the sample rows in string format
+        sample_rows_str = "\n".join(["\t".join(row) for row in sample_rows])
+
+    # in some dialects when there are no rows in the table a
+    # 'ProgrammingError' is returned
+    except ProgrammingError:
+        sample_rows_str = ""
+
+    return (
+        f"{row_count} rows from {table.name} table:\n"
+        f"{columns_str}\n"
+        f"{sample_rows_str}"
+    )
+
+def get_unique_values(engine:Engine, table:Table) -> str:
+    """Gets the unique values of each column in a table using the SQLAlchemy engine"""
+    unique_values = {}
+    for column in table.c:
+        command = select(distinct(column))
+
+        try:
+            # get the sample rows
+            with engine.connect() as connection:
+                result = connection.execute(command)  # type: ignore
+                # shorten values in the sample rows
+                unique_values[column.name] = [str(u) for u in result]
+
+            # save the sample rows in string format
+            # sample_rows_str = "\n".join(["\t".join(row) for row in sample_rows])
+            # in some dialects when there are no rows in the table a
+            # 'ProgrammingError' is returned
+        except ProgrammingError:
+            sample_rows_str = ""
+
+    output_str = f"Unique values of each column in {table.name}: \n"
+    for column, values in unique_values.items():
+        output_str += f"{column} has {len(values)} unique values: {" ".join(values[:20])}"
+        if len(values) > 20:
+            output_str += ", ...."
+        output_str += "\n"
+
+    return output_str
+
+def get_info_sqlalchemy(uri:str) -> dict[str, str] | None:
+    """Gets the dialect name, accessible tables and table schemas using the SQLAlchemy engine"""
+    engine = create_engine(uri)
+    # Get dialect name using inspector
+    inspector = inspect(engine)
+    dialect = inspector.dialect.name
+    # Metadata for tables and columns
+    m = MetaData()
+    m.reflect(engine)
+
+    tables = {}
+    for table in m.tables.values():
+        tables[table.name] = str(CreateTable(table).compile(engine)).rstrip()
+        tables[table.name] += "\n\n/*"
+        tables[table.name] += "\n" + get_sample_rows(engine, table)+"\n"
+        tables[table.name] += "\n" + get_unique_values(engine, table)+"\n"
+        tables[table.name] += "*/"
+
+    return {'sql_dialect': dialect, 'tables': ", ".join(tables.keys()), 'tables_schema': "\n\n".join(tables.values())}
+
 def extract_code_blocks(text):
     pattern = r"```(?:\w+)?\n(.*?)\n```"
     matches = re.findall(pattern, text, re.DOTALL)
     return matches
 
 if __name__ == "__main__":
-    pass
+    from dotenv import load_dotenv
+    import os
+    load_dotenv()
+
+    uri = os.getenv("POSTGRES_URI")
+    print(get_info_sqlalchemy(uri))

var.py

@@ -30,7 +30,7 @@ correctness, efficiency, and security in your SQL queries.\
 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.
+7. **Result**: Don't return the result of the query, return only 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.