Upload 6 files

database_table_descriptions.csv

@@ -0,0 +1,9 @@
+Table,Description
+productlines,"Stores information about the different product lines offered by the company, including a unique name, textual description, HTML description, and image. Categorizes products into different lines."
+products,"Contains details of each product sold by the company, including code, name, product line, scale, vendor, description, stock quantity, buy price, and MSRP. Linked to the productlines table."
+offices,"Holds data on the company's sales offices, including office code, city, phone number, address, state, country, postal code, and territory. Each office is uniquely identified by its office code."
+employees,"Stores information about employees, including number, last name, first name, job title, contact info, and office code. Links to offices and maps organizational structure through the reportsTo attribute."
+customers,"Captures data on customers, including customer number, name, contact details, address, assigned sales rep, and credit limit. Central to managing customer relationships and sales processes."
+payments,"Records payments made by customers, tracking the customer number, check number, payment date, and amount. Linked to the customers table for financial tracking and account management."
+orders,"Details each sales order placed by customers, including order number, dates, status, comments, and customer number. Linked to the customers table, tracking sales transactions."
+orderdetails,"Describes individual line items for each sales order, including order number, product code, quantity, price, and order line number. Links orders to products, detailing the items sold."

examples.py

@@ -0,0 +1,25 @@
+import streamlit as st
+from langchain_community.vectorstores import FAISS
+from langchain_core.example_selectors import SemanticSimilarityExampleSelector
+from langchain_openai import OpenAIEmbeddings
+
+# Example queries remain the same
+examples = [
+    {
+        "input": "Retrieve the user  who have placed the highest total value of orders.",
+        "query": "SELECT u.username, SUM(p.price) AS total_order_value FROM users u JOIN orders o ON u.user_id = o.user_id JOIN products p ON o.product_id = p.product_id GROUP BY u.username ORDER BY total_order_value DESC LIMIT 10;"
+    },
+    # ... (rest of the examples)
+]
+
+@st.cache_resource
+def get_example_selector(api_key):
+    embeddings = OpenAIEmbeddings(api_key=api_key)
+    example_selector = SemanticSimilarityExampleSelector.from_examples(
+        examples,
+        embeddings,
+        FAISS,
+        k=2,
+        input_variables=["input"],
+    )
+    return example_selector

langchain_utils.py

@@ -0,0 +1,71 @@
+import streamlit as st
+from langchain_community.utilities.sql_database import SQLDatabase
+from langchain.chains import create_sql_query_chain
+from langchain_openai import ChatOpenAI
+from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
+from langchain.memory import ChatMessageHistory
+from operator import itemgetter
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from table_details import create_table_chain
+from prompts import create_prompts
+
+def get_db_uri(credentials):
+    return f"postgresql+psycopg2://{credentials['user']}:{credentials['password']}@{credentials['host']}:{credentials['port']}/{credentials['database']}"
+
+@st.cache_resource
+def get_chain(_db_uri, api_key):
+    """Create the langchain with the provided credentials"""
+    try:
+        db = SQLDatabase.from_uri(_db_uri)
+        llm = ChatOpenAI(temperature=0.7, model="gpt-3.5-turbo", api_key=api_key)
+        
+        # Get the table chain and prompts
+        table_chain = create_table_chain(api_key)
+        final_prompt, answer_prompt = create_prompts(api_key)
+        
+        generate_query = create_sql_query_chain(llm, db, final_prompt)
+        execute_query = QuerySQLDataBaseTool(db=db)
+        rephrase_answer = answer_prompt | llm | StrOutputParser()
+        
+        chain = (
+            RunnablePassthrough.assign(table_names_to_use=table_chain) |
+            RunnablePassthrough.assign(query=generate_query).assign(
+                result=itemgetter("query") | execute_query
+            ) | rephrase_answer
+        )
+        
+        return chain
+    except Exception as e:
+        st.error(f"Error creating chain: {str(e)}")
+        return None
+
+def create_history(messages):
+    history = ChatMessageHistory()
+    for message in messages:
+        if message["role"] == "user":
+            history.add_user_message(message["content"])
+        else:
+            history.add_ai_message(message["content"])
+    return history
+
+def invoke_chain(question, messages, db_credentials, api_key):
+    try:
+        db_uri = get_db_uri(db_credentials)
+        chain = get_chain(db_uri, api_key)
+        if chain is None:
+            return "Sorry, I couldn't connect to the database. Please check your credentials."
+        
+        history = create_history(messages)
+        response = chain.invoke({
+            "question": question,
+            "top_k": 100,
+            "messages": history.messages
+        })
+        
+        history.add_user_message(question)
+        history.add_ai_message(response)
+        return response
+        
+    except Exception as e:
+        return f"An error occurred: {str(e)}"

main.py

@@ -0,0 +1,125 @@
+import streamlit as st
+from langchain_openai import ChatOpenAI
+from langchain_utils import invoke_chain
+from langchain_community.callbacks import get_openai_callback
+
+def init_session_state():
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    if "connected" not in st.session_state:
+        st.session_state.connected = False
+
+def create_sidebar():
+    with st.sidebar:
+        st.title("Postgres Credentials")
+        st.subheader("Enter your Credentials & Connect")
+        
+        # Database credentials
+        host = st.text_input("Host", value="localhost")
+        port = st.text_input("Port", value="5432")
+        user = st.text_input("User", value="postgres")
+        password = st.text_input("Password", type="password")
+        database = st.text_input("Database")
+        
+        # OpenAI API key
+        api_key = st.text_input("OpenAI API Key", type="password")
+        
+        # Connect button
+        if st.button("Connect", use_container_width=True):
+            try:
+                # Store credentials in session state
+                st.session_state.db_credentials = {
+                    "host": host,
+                    "port": port,
+                    "user": user,
+                    "password": password,
+                    "database": database
+                }
+                st.session_state.api_key = api_key
+                st.session_state.connected = True
+                st.success("Successfully connected!")
+            except Exception as e:
+                st.error(f"Connection failed: {str(e)}")
+                st.session_state.connected = False
+
+def main():
+    init_session_state()
+    create_sidebar()
+    
+    st.title("Chat with Postgres DB")
+    
+    if not st.session_state.connected:
+        st.info("Please enter your credentials in the sidebar and connect first.")
+        return
+    
+    # Display the welcome message with the database icon using markdown
+    st.markdown("""
+        <div style='display: flex; align-items: center; gap: 10px;'>
+            <span style='font-size: 24px;'>🗄️</span>
+            <span>Hello! I'm a QualityKiosk's SQL assistant. Ask me anything about your database.</span>
+        </div>
+    """, unsafe_allow_html=True)
+    
+    # Display chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+    
+    # Chat input
+    if prompt := st.chat_input("Type a message..."):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                with get_openai_callback() as cb:
+                    response = invoke_chain(
+                        prompt,
+                        st.session_state.messages,
+                        st.session_state.db_credentials,
+                        st.session_state.api_key
+                    )
+                    print(f"OpenAI Stats: {cb}")
+                st.markdown(response)
+        
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
+if __name__ == "__main__":
+    # Set dark theme and wide layout
+    st.set_page_config(
+        page_title="Chat with Postgres DB",
+        layout="wide",
+        initial_sidebar_state="expanded",
+        # Optional: Add a custom theme
+        menu_items={
+            'Get Help': 'https://www.qualitykiosk.com',
+            'About': "# Chat with Postgres DB\nA QualityKiosk's SQL Assistant"
+        }
+    )
+    
+    # Add custom CSS for dark theme and styling
+    st.markdown("""
+        <style>
+            .stApp {
+                background-color: #1E1E1E;
+                color: #FFFFFF;
+            }
+            .stSidebar {
+                background-color: #262626;
+            }
+            .stButton>button {
+                background-color: #0E86D4;
+                color: white;
+            }
+            .stTextInput>div>div>input {
+                background-color: #333333;
+                color: white;
+            }
+            .stMarkdown {
+                color: white;
+            }
+        </style>
+    """, unsafe_allow_html=True)
+    
+    main()

prompts.py

@@ -0,0 +1,39 @@
+from examples import get_example_selector
+from langchain_core.prompts import (
+    ChatPromptTemplate,
+    MessagesPlaceholder,
+    FewShotChatMessagePromptTemplate,
+    PromptTemplate
+)
+
+example_prompt = ChatPromptTemplate.from_messages([
+    ("human", "{input}\nSQLQuery:"),
+    ("ai", "{query}"),
+])
+
+def create_prompts(api_key):
+    example_selector = get_example_selector(api_key)
+    
+    few_shot_prompt = FewShotChatMessagePromptTemplate(
+        example_prompt=example_prompt,
+        example_selector=example_selector,
+        input_variables=["input", "top_k"],
+    )
+
+    final_prompt = ChatPromptTemplate.from_messages([
+        ("system", "You are a Postgres SQL expert. Given an input question, create a syntactically correct Postgres SQL query to run. Unless otherwise specified.\n\nHere is the relevant table info: {table_info}\n\nBelow are a number of examples of questions and their corresponding SQL queries."),
+        few_shot_prompt,
+        MessagesPlaceholder(variable_name="messages"),
+        ("human", "{input}"),
+    ])
+
+    answer_prompt = PromptTemplate.from_template(
+        """Given the following user question, corresponding SQL query, and SQL result, answer the user question.
+
+    Question: {question}
+    SQL Query: {query}
+    SQL Result: {result}
+    Answer: """
+    )
+
+    return final_prompt, answer_prompt

table_details.py

@@ -0,0 +1,42 @@
+import pandas as pd
+import streamlit as st
+from operator import itemgetter
+from langchain.chains.openai_tools import create_extraction_chain_pydantic
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_openai import ChatOpenAI
+from typing import List
+
+class Table(BaseModel):
+    """Table in SQL database."""
+    name: str = Field(description="Name of table in SQL database.")
+
+def get_tables(tables: List[Table]) -> List[str]:
+    return [table.name for table in tables]
+
+@st.cache_data
+def get_table_details():
+    try:
+        table_description = pd.read_excel("database_table_descriptions.xlsx")
+        table_details = ""
+        for index, row in table_description.iterrows():
+            table_details += f"Table Name:{row['Table']}\nTable Description:{row['Description']}\n\n"
+        return table_details
+    except Exception as e:
+        st.error(f"Error reading table descriptions: {str(e)}")
+        return ""
+
+def create_table_chain(api_key):
+    table_details = get_table_details()
+    table_details_prompt = f"""Return the names of ALL the SQL tables that MIGHT be relevant to the user question. \
+    The tables are:
+
+    {table_details}
+
+    Remember to include ALL POTENTIALLY RELEVANT tables, even if you're not sure that they're needed."""
+
+    llm = ChatOpenAI(temperature=0.7, model="gpt-3.5-turbo", api_key=api_key)
+    return (
+        {"input": itemgetter("question")} | 
+        create_extraction_chain_pydantic(Table, llm, system_message=table_details_prompt) | 
+        get_tables
+    )