Update app.py

app.py

@@ -4,7 +4,7 @@ import streamlit as st
 from werkzeug.security import generate_password_hash, check_password_hash
 from langchain_groq import ChatGroq
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_core.documents import Document
+from langchain_community.document_loaders.csv_loader import CSVLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_core.vectorstores import InMemoryVectorStore
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
@@ -12,72 +12,46 @@ from langchain.tools import Tool
 from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain_core.messages import HumanMessage, AIMessage
 
-# ----------------- Database Setup -----------------
+# --- Database Setup ---
 @st.cache_resource
 def init_db():
     conn = sqlite3.connect('users.db', check_same_thread=False)
     c = conn.cursor()
     
-    # Users Table
+    # Create users table with admin flag
     c.execute('''CREATE TABLE IF NOT EXISTS users
                  (id INTEGER PRIMARY KEY AUTOINCREMENT,
                   username TEXT UNIQUE NOT NULL,
                   password TEXT NOT NULL,
                   previous_chat_history TEXT,
                   previous_products_bought TEXT,
-                  is_admin BOOLEAN DEFAULT 0)''')
+                  is_admin INTEGER DEFAULT 0)''')
     
-    # Products Table
-    c.execute('''CREATE TABLE IF NOT EXISTS products
-                 (product_id INTEGER PRIMARY KEY AUTOINCREMENT,
-                  name TEXT NOT NULL,
-                  description TEXT,
-                  price REAL,
-                  stock INTEGER,
-                  category TEXT,
-                  features TEXT)''')
-    
-    # Company Settings Table
+    # Create company settings table
     c.execute('''CREATE TABLE IF NOT EXISTS company_settings
-                 (id INTEGER PRIMARY KEY,
-                  company_name TEXT DEFAULT 'TechElectronics',
-                  business_type TEXT DEFAULT 'Consumer Electronics Retailer',
-                  key_features TEXT DEFAULT 'Cutting-edge technology, Competitive pricing, Excellent customer service',
-                  payment_link TEXT DEFAULT 'https://www.example.com/payment')''')
+                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
+                  company_name TEXT,
+                  company_business TEXT,
+                  agent_name TEXT,
+                  key_features TEXT)''')
     
-    # Initialize Default Data
+    # Initialize default company settings if empty
     c.execute('SELECT COUNT(*) FROM company_settings')
     if c.fetchone()[0] == 0:
         c.execute('''INSERT INTO company_settings 
-                     (company_name, business_type, key_features, payment_link)
+                     (company_name, company_business, agent_name, key_features)
                      VALUES (?, ?, ?, ?)''',
-                 ('TechElectronics', 
-                  'Consumer Electronics Retailer',
-                  'Cutting-edge technology, Competitive pricing, Excellent customer service',
-                  'https://www.example.com/payment'))
-    
-    c.execute('SELECT COUNT(*) FROM products')
-    if c.fetchone()[0] == 0:
-        sample_products = [
-            ('Smartphone X', 'Flagship smartphone with AI camera', 799.99, 50, 
-             'Phones', '5G, 128GB Storage, 48MP Camera'),
-            ('4K Smart TV', '55-inch 4K Android TV', 899.99, 30, 
-             'TVs', 'HDR10+, Dolby Atmos, Smart Assistant')
-        ]
-        c.executemany('''INSERT INTO products 
-                      (name, description, price, stock, category, features)
-                      VALUES (?, ?, ?, ?, ?, ?)''', sample_products)
+                  ("TechElectronics", "Consumer Electronics Retailer", "Alex",
+                   "Cutting-edge technology, Competitive pricing, Excellent customer service"))
     
     conn.commit()
     return conn
 
 conn = init_db()
 
-# ----------------- Data Management Classes -----------------
-# User Mangement system
+# --- User Class ---
 class User:
-    def __init__(self, id, username, password, chat_history=None, 
-                 products_bought=None, is_admin=False):
+    def __init__(self, id, username, password, chat_history=None, products_bought=None, is_admin=False):
         self.id = id
         self.username = username
         self.password = password
@@ -88,147 +62,117 @@ class User:
     @classmethod
     def create(cls, username, password, is_admin=False):
         hashed_pw = generate_password_hash(password)
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            c.execute('''INSERT INTO users (username, password, is_admin)
-                          VALUES (?, ?, ?)''',
-                     (username, hashed_pw, is_admin))
-            return cls(c.lastrowid, username, hashed_pw, is_admin=is_admin)
+        conn = sqlite3.connect('users.db')
+        c = conn.cursor()
+        c.execute('INSERT INTO users (username, password, is_admin) VALUES (?, ?, ?)',
+                  (username, hashed_pw, is_admin))
+        user_id = c.lastrowid
+        conn.commit()
+        conn.close()
+        return cls(user_id, username, hashed_pw, is_admin=is_admin)
 
     @classmethod
     def get_by_username(cls, username):
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            c.execute('SELECT * FROM users WHERE username = ?', (username,))
-            user = c.fetchone()
-            if user:
-                return cls(user[0], user[1], user[2], 
-                         eval(user[3]) if user[3] else [],
-                         eval(user[4]) if user[4] else [],
-                         user[5])
+        conn = sqlite3.connect('users.db')
+        c = conn.cursor()
+        c.execute('SELECT * FROM users WHERE username = ?', (username,))
+        user = c.fetchone()
+        conn.close()
+        if user:
+            return cls(user[0], user[1], user[2], 
+                      eval(user[3]) if user[3] else [],
+                      eval(user[4]) if user[4] else [],
+                      is_admin=user[5])
         return None
 
     def update_chat_history(self, new_messages):
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            updated_history = self.chat_history + new_messages
-            c.execute('''UPDATE users SET previous_chat_history = ?
-                      WHERE id = ?''', (str(updated_history), self.id))
-            conn.commit()
+        conn = sqlite3.connect('users.db')
+        c = conn.cursor()
+        updated_history = self.chat_history + new_messages
+        c.execute('UPDATE users SET previous_chat_history = ? WHERE id = ?',
+                 (str(updated_history), self.id))
+        conn.commit()
+        conn.close()
 
     def update_products_bought(self, new_products):
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            updated_products = self.products_bought + new_products
-            c.execute('''UPDATE users SET previous_products_bought = ?
-                      WHERE id = ?''', (str(updated_products), self.id))
-            conn.commit()
+        conn = sqlite3.connect('users.db')
+        c = conn.cursor()
+        updated_products = self.products_bought + new_products
+        c.execute('UPDATE users SET previous_products_bought = ? WHERE id = ?',
+                 (str(updated_products), self.id))
+        conn.commit()
+        conn.close()
 
-# Product Management system
-class ProductManager:
-    @staticmethod
-    def get_all_products():
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            c.execute('SELECT * FROM products')
-            return c.fetchall()
+# --- Company Settings Functions ---
+def get_company_settings():
+    conn = sqlite3.connect('users.db')
+    c = conn.cursor()
+    c.execute('SELECT * FROM company_settings ORDER BY id DESC LIMIT 1')
+    settings = c.fetchone()
+    conn.close()
+    if settings:
+        return {
+            'company_name': settings[1],
+            'company_business': settings[2],
+            'agent_name': settings[3],
+            'key_features': settings[4]
+        }
+    return None
 
-    @staticmethod
-    def update_product(product_id, name, description, price, stock, category, features):
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            c.execute('''UPDATE products SET
-                       name=?, description=?, price=?, stock=?,
-                       category=?, features=?
-                       WHERE product_id=?''',
-                    (name, description, price, stock, category, features, product_id))
-            conn.commit()
+def update_company_settings(company_name, company_business, agent_name, key_features):
+    conn = sqlite3.connect('users.db')
+    c = conn.cursor()
+    c.execute('''INSERT INTO company_settings 
+                 (company_name, company_business, agent_name, key_features)
+                 VALUES (?, ?, ?, ?)''',
+              (company_name, company_business, agent_name, key_features))
+    conn.commit()
+    conn.close()
 
-    @staticmethod
-    def add_product(name, description, price, stock, category, features):
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            c.execute('''INSERT INTO products 
-                      (name, description, price, stock, category, features)
-                      VALUES (?, ?, ?, ?, ?, ?)''',
-                    (name, description, price, stock, category, features))
-            conn.commit()
+# --- AI Agent Setup ---
+os.environ["GROQ_API_KEY"] = st.secrets["GROQ_API_KEY"]
+llm = ChatGroq(
+    temperature=0.1,
+    model_name="llama3-8b-8192",
+    api_key=st.secrets["GROQ_API_KEY"],
+)
 
-    @staticmethod
-    def delete_product(product_id):
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            c.execute('DELETE FROM products WHERE product_id=?', (product_id,))
-            conn.commit()
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-# Company Setting
-class CompanyManager:
-    @staticmethod
-    def get_settings():
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            c.execute('SELECT * FROM company_settings WHERE id=1')
-            settings = c.fetchone()
-            return {
-                'company_name': settings[1],
-                'business_type': settings[2],
-                'key_features': settings[3],
-                'payment_link': settings[4]
-            }
+@st.cache_resource
+def load_data():
+    loader = CSVLoader(file_path="electronics_products.csv")
+    docs = loader.load()
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
+    splits = text_splitter.split_documents(docs)
+    vectorstore = InMemoryVectorStore.from_documents(documents=splits, embedding=embeddings)
+    return vectorstore.as_retriever()
 
-    @staticmethod
-    def update_settings(name, business, features, payment_link):
-        with sqlite3.connect('users.db') as conn:
-            c = conn.cursor()
-            c.execute('''UPDATE company_settings SET
-                       company_name=?, business_type=?, 
-                       key_features=?, payment_link=?
-                       WHERE id=1''',
-                    (name, business, features, payment_link))
-            conn.commit()
+retriever = load_data()
 
-# ----------------- AI Agent Setup -----------------
-@st.cache_resource
-def initialize_ai():
-    os.environ["GROQ_API_KEY"] = st.secrets["GROQ_API_KEY"]
-    llm = ChatGroq(
-        temperature=0.1,
-        model_name="llama3-8b-8192",
-        api_key=st.secrets["GROQ_API_KEY"],
-    )
-    
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    
-    def refresh_vectorstore():
-        products = ProductManager.get_all_products()
-        docs = []
-        for p in products:
-            docs.append(Document(
-                page_content=f"Name: {p[1]}\nDescription: {p[2]}\nPrice: {p[3]}\nStock: {p[4]}\nCategory: {p[5]}\nFeatures: {p[6]}",
-                metadata={"product_id": p[0]}
-            ))
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
-        splits = text_splitter.split_documents(docs)
-        return InMemoryVectorStore.from_documents(splits, embeddings).as_retriever()
-    
-    return llm, embeddings, refresh_vectorstore
+def retrieve_query(query: str):
+    return retriever.get_relevant_documents(query)
 
-llm, embeddings, refresh_vectorstore = initialize_ai()
+tool = Tool(
+    name="retriever",
+    func=retrieve_query,
+    description="Useful for retrieving product information"
+)
 
-def get_system_prompt():
-    settings = CompanyManager.get_settings()
-    return f"""
-    You are {settings['company_name']}'s AI Sales Assistant specializing in {settings['business_type']}.
-    Key Features: {settings['key_features']}
-    
-    Guidelines:
-    1. Only recommend products with stock > 0
-    2. Always verify prices before confirming
-    3. Use payment link: {settings['payment_link']}
-    4. Maintain professional and friendly tone
+@st.cache_resource
+def create_agent_executor(company_name, company_business, agent_name, key_features):
+    system_prompt = """
+    You are {agent_name}, the AI Sales Assistant for {company_name} ({company_business}).
+
+    Company Profile:
+    - Company Name: {company_name}
+    - Business: {company_business}
+    - Key Features: {key_features}
 
-    Current Products:
-    {[p[1] for p in ProductManager.get_all_products()]}
+    Product Availability:
+    - Only recommend and discuss products listed in the provided document.
+    - Do not suggest unavailable or out-of-stock products.
+    - Always verify product availability before making recommendations.
 
     Conversation Flow:
     1. Introduction
@@ -243,131 +187,71 @@ def get_system_prompt():
     - Maintain natural, professional conversations
     - Follow company policies
     - Be helpful and polite
+    - Always cross-check product recommendations with the available inventory
     """
 
-def create_agent_executor():
-    retriever = refresh_vectorstore()
-    
-    tool = Tool(
-        name="product_retriever",
-        func=lambda q: retriever.get_relevant_documents(q),
-        description="Access current product information"
+    formatted_system_prompt = system_prompt.format(
+        agent_name=agent_name,
+        company_name=company_name,
+        company_business=company_business,
+        key_features=key_features
     )
-    
+
     prompt = ChatPromptTemplate.from_messages([
-        ("system", get_system_prompt()),
+        ("system", formatted_system_prompt),
         MessagesPlaceholder(variable_name="chat_history"),
         ("human", "{input}"),
         MessagesPlaceholder(variable_name="agent_scratchpad")
     ])
-    
+
     agent = create_tool_calling_agent(llm, [tool], prompt)
     return AgentExecutor(agent=agent, tools=[tool], verbose=True)
 
-# ----------------- Streamlit UI -----------------
+# --- Admin Dashboard ---
 def admin_dashboard():
-    st.sidebar.header("Company Configuration")
-    with st.sidebar.form("Company Settings"):
-        settings = CompanyManager.get_settings()
-        company_name = st.text_input("Company Name", value=settings['company_name'])
-        business_type = st.text_area("Business Type", value=settings['business_type'])
-        key_features = st.text_area("Key Features", value=settings['key_features'])
-        payment_link = st.text_input("Payment Link", value=settings['payment_link'])
-        
-        if st.form_submit_button("Update Settings"):
-            CompanyManager.update_settings(company_name, business_type, 
-                                        key_features, payment_link)
-            st.success("Settings updated!")
-            st.cache_resource.clear()
-
-    st.header("Product Management")
-    tab1, tab2, tab3 = st.tabs(["Add Product", "Edit Products", "Delete Product"])
-    
-    with tab1:
-        with st.form("Add Product"):
-            name = st.text_input("Product Name")
-            desc = st.text_area("Description")
-            price = st.number_input("Price", min_value=0.0)
-            stock = st.number_input("Stock", min_value=0)
-            category = st.text_input("Category")
-            features = st.text_input("Features")
-            if st.form_submit_button("Add Product"):
-                ProductManager.add_product(name, desc, price, stock, category, features)
-                refresh_vectorstore()
-                st.success("Product added!")
-
-    with tab2:
-        products = ProductManager.get_all_products()
-        selected = st.selectbox("Select Product", products, format_func=lambda x: x[1])
-        if selected:
-            with st.form("Edit Product"):
-                name = st.text_input("Name", value=selected[1])
-                desc = st.text_area("Description", value=selected[2])
-                price = st.number_input("Price", value=selected[3])
-                stock = st.number_input("Stock", value=selected[4])
-                category = st.text_input("Category", value=selected[5])
-                features = st.text_input("Features", value=selected[6])
-                if st.form_submit_button("Update"):
-                    ProductManager.update_product(selected[0], name, desc, 
-                                               price, stock, category, features)
-                    refresh_vectorstore()
-                    st.success("Product updated!")
-
-    with tab3:
-        products = ProductManager.get_all_products()
-        selected = st.selectbox("Select to Delete", products, format_func=lambda x: x[1])
-        if st.button("Delete"):
-            ProductManager.delete_product(selected[0])
-            refresh_vectorstore()
-            st.success("Product deleted!")
-
-def user_interface():
-    settings = CompanyManager.get_settings()
-    st.title(f"{settings['company_name']} AI Assistant")
+    st.header("Admin Dashboard")
     
-    if 'chat_history' not in st.session_state:
-        st.session_state.chat_history = st.session_state.user.chat_history
-    
-    for msg in st.session_state.chat_history:
-        role = "user" if msg["type"] == "human" else "assistant"
-        with st.chat_message(role):
-            st.write(msg["content"])
+    # Company Settings Form
+    settings = get_company_settings()
+    with st.expander("Company Configuration", expanded=True):
+        with st.form("company_settings"):
+            st.subheader("Company Settings")
+            new_company_name = st.text_input("Company Name", value=settings['company_name'])
+            new_company_business = st.text_input("Company Business", value=settings['company_business'])
+            new_agent_name = st.text_input("Agent Name", value=settings['agent_name'])
+            new_key_features = st.text_area("Key Features", value=settings['key_features'])
+            
+            if st.form_submit_button("Update Settings"):
+                update_company_settings(new_company_name, new_company_business, 
+                                      new_agent_name, new_key_features)
+                st.success("Company settings updated successfully!")
 
-    if prompt := st.chat_input("How can I help you today?"):
-        st.chat_message("user").write(prompt)
-        
-        agent_executor = create_agent_executor()
-        response = agent_executor.invoke({
-            "input": prompt,
-            "chat_history": [
-                HumanMessage(content=msg["content"]) if msg["type"] == "human" 
-                else AIMessage(content=msg["content"]) 
-                for msg in st.session_state.chat_history
-            ]
-        })["output"]
-        
-        with st.chat_message("assistant"):
-            st.write(response)
-            if settings['payment_link'] in response:
-                st.markdown(f"[Complete Payment Here]({settings['payment_link']})")
-
-        new_messages = [
-            {"type": "human", "content": prompt},
-            {"type": "ai", "content": response}
-        ]
-        st.session_state.user.update_chat_history(new_messages)
-        st.session_state.chat_history += new_messages
+    # Product Management
+    with st.expander("Product Management"):
+        st.subheader("Update Product Catalog")
+        uploaded_file = st.file_uploader("Upload new product CSV", type="csv")
+        if uploaded_file is not None:
+            # Save the uploaded file
+            with open("electronics_products.csv", "wb") as f:
+                f.write(uploaded_file.getbuffer())
+            # Clear cache to reload data
+            load_data.clear()
+            st.success("Product catalog updated successfully!")
 
+# --- Main Application ---
 def main():
-    st.set_page_config(page_title="AI Sales Assistant", layout="wide")
+    st.title("AI Sales Assistant 🤖")
     
     if 'user' not in st.session_state:
         st.session_state.user = None
-    
+        st.session_state.chat_history = []
+
+    # Authentication
     if not st.session_state.user:
-        auth_col1, auth_col2 = st.columns(2)
+        st.header("Login/Register")
+        tab1, tab2 = st.tabs(["Login", "Register"])
         
-        with auth_col1:
+        with tab1:
             with st.form("Login"):
                 username = st.text_input("Username")
                 password = st.text_input("Password", type="password")
@@ -375,25 +259,63 @@ def main():
                     user = User.get_by_username(username)
                     if user and check_password_hash(user.password, password):
                         st.session_state.user = user
+                        st.session_state.chat_history = user.chat_history
                         st.rerun()
                     else:
                         st.error("Invalid credentials")
         
-        with auth_col2:
+        with tab2:
             with st.form("Register"):
                 new_user = st.text_input("New Username")
                 new_pass = st.text_input("New Password", type="password")
-                if st.form_submit_button("Create Account"):
+                admin_key = st.text_input("Admin Key (optional)", type="password")
+                if st.form_submit_button("Register"):
                     if User.get_by_username(new_user):
-                        st.error("Username exists")
+                        st.error("Username already exists")
                     else:
-                        st.session_state.user = User.create(new_user, new_pass)
+                        is_admin = admin_key == st.secrets.get("ADMIN_KEY", "")
+                        user = User.create(new_user, new_pass, is_admin)
+                        st.session_state.user = user
+                        st.session_state.chat_history = []
                         st.rerun()
     else:
         if st.session_state.user.is_admin:
             admin_dashboard()
         else:
-            user_interface()
+            settings = get_company_settings()
+            agent_executor = create_agent_executor(
+                settings['company_name'],
+                settings['company_business'],
+                settings['agent_name'],
+                settings['key_features']
+            )
+            
+            st.header(f"Welcome to {settings['company_name']}, {st.session_state.user.username}!")
+            st.subheader(f"Chat with {settings['agent_name']}, your AI Sales Assistant")
+            
+            if prompt := st.chat_input("Type your message here..."):
+                with st.chat_message("user"):
+                    st.write(prompt)
+
+                with st.chat_message("assistant"):
+                    response = agent_executor.invoke({
+                        "input": prompt,
+                        "chat_history": [HumanMessage(content=msg["content"]) if msg["type"] == "human" 
+                                      else AIMessage(content=msg["content"]) 
+                                      for msg in st.session_state.chat_history]
+                    })["output"]
+                    st.write(response)
+                    
+                    if "https://www.example.com/payment" in response:
+                        st.session_state.user.update_products_bought(["Latest Product"])
+                        st.success("Product added to your purchases!")
+
+                new_messages = [
+                    {"type": "human", "content": prompt},
+                    {"type": "ai", "content": response}
+                ]
+                st.session_state.user.update_chat_history(new_messages)
+                st.session_state.chat_history += new_messages
 
 if __name__ == "__main__":
     main()