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_community.document_loaders.csv_loader import CSVLoader
+from langchain_core.documents import Document
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_core.vectorstores import InMemoryVectorStore
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
@@ -12,253 +12,388 @@ from langchain.tools import Tool
 from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain_core.messages import HumanMessage, AIMessage
 
-
-# --- Database Setup ---
-
-# Database initialization with caching
+# ----------------- Database Setup -----------------
 @st.cache_resource
 def init_db():
     conn = sqlite3.connect('users.db', check_same_thread=False)
     c = conn.cursor()
+    
+    # Users Table
     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)''')
+                  previous_products_bought TEXT,
+                  is_admin BOOLEAN 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
+    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')''')
+    
+    # Initialize Default Data
+    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)
+                     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)
+    
     conn.commit()
     return conn
 
 conn = init_db()
 
-
+# ----------------- Data Management Classes -----------------
+# User Mangement system
 class User:
-    def __init__(self, id, username, password, chat_history = None, products_bought = None):
+    def __init__(self, id, username, password, chat_history=None, 
+                 products_bought=None, is_admin=False):
         self.id = id
         self.username = username
         self.password = password
         self.chat_history = chat_history or []
         self.products_bought = products_bought or []
+        self.is_admin = is_admin
 
-    # To register a new user
     @classmethod
-    def create(cls, username, password):
+    def create(cls, username, password, is_admin=False):
         hashed_pw = generate_password_hash(password)
-        conn = sqlite3.connect('users.db')
-        c = conn.cursor()
-        c.execute('INSERT INTO users (username, password) VALUES (?, ?)',(username, hashed_pw))
-        user_id = c.lastrowid
-        conn.commit()
-        conn.close()
-        return cls(user_id, username, hashed_pw)    
+        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)
 
-     # To retrieve an existing user from the database by username.
     @classmethod
     def get_by_username(cls, username):
-        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 [])
+        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])
         return None
 
     def update_chat_history(self, new_messages):
-        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()    
+        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()
 
     def update_products_bought(self, new_products):
-        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()    
-        
+        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()
+
+# 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()
+
+    @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()
+
+    @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 ---
-# Load the LLM model from Groq
-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()
 
-# Load the HuggingFace Embeddings
-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]
+            }
 
-# Load and process CSV data
+    @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()
+
+# ----------------- AI Agent Setup -----------------
 @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()
+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
 
-retriever = load_data()
+llm, embeddings, refresh_vectorstore = initialize_ai()
 
-def retrieve_query(query: str):
-    """Retrieves documents related to the query."""
-    return retriever.get_relevant_documents(query)
+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
 
-tool = Tool(
-    name="retriever",
-    func=retrieve_query,
-    description="Useful for retrieving product information"
-)
+    Current Products:
+    {[p[1] for p in ProductManager.get_all_products()]}
 
-# Agent setup
-# System prompt template
-system_prompt = """
-      You are {agent_name}, the AI Sales Assistant for {company_name} ({company_business}).
+    Conversation Flow:
+    1. Introduction
+    2. Qualification
+    3. Understanding Needs
+    4. Needs Analysis
+    5. Solution Presentation (Only recommend available products)
+    6. Confirmation
+    7. If the prospect agrees to purchase, thank them and provide the payment link: https://www.example.com/payment
 
-      Company Profile:
-      - Company Name: {company_name}
-      - Business: {company_business}
-      - Key Features: {key_features}
+    Guidelines:
+    - Maintain natural, professional conversations
+    - Follow company policies
+    - Be helpful and polite
+    """
 
-      Product Availability:
-      - Only recommend and discuss products listed in the provided in document.
-      - Do not suggest unavailable or out-of-stock products.
-      - Always verify product availability before making recommendations.
+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"
+    )
+    
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", get_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)
 
-      Conversation Flow:
-      1. Introduction
-      2. Qualification
-      3. Understanding Needs
-      4. Needs Analysis
-      5. Solution Presentation (Only recommend available products)
-      6. Confirmation
-      7. If the prospect agrees to purchase, thank them and provide the payment link: https://www.example.com/payment
+# ----------------- Streamlit UI -----------------
+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()
 
-      Guidelines:
-      - Maintain natural, professional conversations
-      - Follow company policies
-      - Be helpful and polite
-      - Always cross-check product recommendations with the available inventory in the CSV document
-      """
+    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!")
 
-# Define the company and agent details
-company_name = "TechElectronics"
-company_business = "Consumer Electronics Retailer"
-agent_name = "Alex"
-key_features = "Cutting-edge technology, Competitive pricing, Excellent customer service"
+    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!")
 
-# Format the system prompt with the company and agent details
-formatted_system_prompt = system_prompt.format(
-    agent_name=agent_name,
-    company_name=company_name,
-    company_business=company_business,
-    key_features=key_features
-)
+    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!")
 
-prompt = ChatPromptTemplate.from_messages([
-    ("system", formatted_system_prompt),
-    MessagesPlaceholder(variable_name="chat_history"),
-    ("human", "{input}"),
-    MessagesPlaceholder(variable_name="agent_scratchpad")
-])
+def user_interface():
+    settings = CompanyManager.get_settings()
+    st.title(f"{settings['company_name']} AI Assistant")
+    
+    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"])
 
-tools = [tool]
-agent = create_tool_calling_agent(llm, tools, prompt)
-agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+    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
 
-# --- Streamlit UI ---
 def main():
-    st.title(f"{company_name} AI Sales Assistant 🤖")
+    st.set_page_config(page_title="AI Sales Assistant", layout="wide")
     
-    # Initialize session state
     if 'user' not in st.session_state:
         st.session_state.user = None
-        st.session_state.chat_history = []
     
-    # Authentication
     if not st.session_state.user:
-        st.header("Login/Register")
-        tab1, tab2 = st.tabs(["Login", "Register"])
+        auth_col1, auth_col2 = st.columns(2)
         
-        with tab1:
+        with auth_col1:
             with st.form("Login"):
                 username = st.text_input("Username")
                 password = st.text_input("Password", type="password")
                 if st.form_submit_button("Login"):
                     user = User.get_by_username(username)
                     if user and check_password_hash(user.password, password):
-                        # If valid, the user is stored in st.session_state.user and their chat history is loaded.
                         st.session_state.user = user
-                        st.session_state.chat_history = user.chat_history
                         st.rerun()
                     else:
                         st.error("Invalid credentials")
         
-        with tab2:
+        with auth_col2:
             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("Register"):
+                if st.form_submit_button("Create Account"):
                     if User.get_by_username(new_user):
-                        st.error("Username already exists")
+                        st.error("Username exists")
                     else:
-                        user = User.create(new_user, new_pass)
-                        st.session_state.user = user
-                        st.session_state.chat_history = []
+                        st.session_state.user = User.create(new_user, new_pass)
                         st.rerun()
-
     else:
-        # Chat Interface
-        st.header(f"Welcome to {company_name}, {st.session_state.user.username}😊!")
-        st.subheader("Chat with our AI Sales Assistant")
-        
-        # # Display chat history
-        # for msg in st.session_state.chat_history:
-        #     if msg["type"] == "human":
-        #         with st.chat_message("user"):
-        #             st.write(msg["content"])
-        #     else:
-        #         with st.chat_message("assistant"):
-        #             st.write(msg["content"])
-
-        if prompt := st.chat_input("Type you Message here..."):
-            #Add user message to chat
-            with st.chat_message("user"):
-                st.write(prompt)
-
-            # Get AI response
-            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)
-                
-                # Check if payment link is provided
-                if "https://www.example.com/payment" in response:
-                    st.session_state.user.update_products_bought(["Latest Product"])
-                    st.success("Product added to your purchases!")    
-                    
-            # Update chat history in database
-            new_messages = [
-                {"type": "human", "content": prompt},
-                {"type": "ai", "content": response}
-            ]
-            # Both the user’s message and the assistant’s reply are appended to the persistent chat history 
-            # (both in session and in the database), ensuring conversation continuity.
-            st.session_state.user.update_chat_history(new_messages)
-            st.session_state.chat_history += new_messages
+        if st.session_state.user.is_admin:
+            admin_dashboard()
+        else:
+            user_interface()
 
 if __name__ == "__main__":
-    main()    
+    main()