app.py

import os
import sqlite3
import streamlit as st
import re
from typing import List, Dict, Any, Optional
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_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.tools import Tool
from langchain_core.tools import StructuredTool
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.messages import HumanMessage, AIMessage
from langchain.docstore.document import Document

# --- 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)''')
    
    # Company settings table
    c.execute('''CREATE TABLE IF NOT EXISTS company_settings
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  name TEXT NOT NULL,
                  business TEXT NOT NULL,
                  agent_name TEXT NOT NULL,
                  key_features TEXT NOT NULL)''')
    
    # Products table with inventory
    c.execute('''CREATE TABLE IF NOT EXISTS products
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  name TEXT NOT NULL,
                  category TEXT NOT NULL,
                  price REAL NOT NULL,
                  description TEXT NOT NULL,
                  features TEXT NOT NULL,
                  stock INTEGER NOT NULL DEFAULT 0)''')
    
    # Check and update schema if needed
    c.execute("PRAGMA table_info(products)")
    columns = [column[1] for column in c.fetchall()]
    if 'stock' not in columns:
        c.execute('ALTER TABLE products ADD COLUMN stock INTEGER NOT NULL DEFAULT 0')
    
    # Insert default company settings if empty
    c.execute('SELECT COUNT(*) FROM company_settings')
    if c.fetchone()[0] == 0:
        c.execute('''INSERT INTO company_settings 
                     (name, business, agent_name, key_features)
                     VALUES (?, ?, ?, ?)''',
                  ('TechElectronics', 
                   'Consumer Electronics Retailer',
                   'Alex',
                   'Cutting-edge technology, Competitive pricing, Excellent customer service'))
    
    conn.commit()
    return conn

conn = init_db()

# --- Admin Classes ---
class Company:
    @staticmethod
    def get_settings():
        c = conn.cursor()
        c.execute('SELECT * FROM company_settings LIMIT 1')
        return c.fetchone()

    @staticmethod
    def update_settings(name, business, agent_name, key_features):
        c = conn.cursor()
        c.execute('''UPDATE company_settings 
                     SET name=?, business=?, agent_name=?, key_features=?
                     WHERE id=1''',
                   (name, business, agent_name, key_features))
        conn.commit()

class Product:
    @staticmethod
    def get_all():
        c = conn.cursor()
        c.execute('SELECT * FROM products')
        return c.fetchall()
    
    @staticmethod
    def get_in_stock():
        c = conn.cursor()
        c.execute('SELECT * FROM products WHERE stock > 0')
        return c.fetchall()

    @staticmethod
    def get_by_id(product_id):
        c = conn.cursor()
        c.execute('SELECT * FROM products WHERE id = ?', (product_id,))
        return c.fetchone()

    @staticmethod
    def add(name, category, price, description, features, stock):
        c = conn.cursor()
        c.execute('''INSERT INTO products 
                     (name, category, price, description, features, stock)
                     VALUES (?, ?, ?, ?, ?, ?)''',
                  (name, category, price, description, features, stock))
        conn.commit()

    @staticmethod
    def delete(product_id):
        c = conn.cursor()
        c.execute('DELETE FROM products WHERE id=?', (product_id,))
        conn.commit()

    @staticmethod
    def update_stock(product_id, new_stock):
        c = conn.cursor()
        c.execute('UPDATE products SET stock=? WHERE id=?', (new_stock, product_id))
        conn.commit()

# --- User Class ---
class User:
    def __init__(self, id, username, password, chat_history=None, products_bought=None):
        self.id = id
        self.username = username
        self.password = password
        self.chat_history = chat_history or []
        self.products_bought = products_bought or []

    @classmethod
    def create(cls, username, password):
        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)

    @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 [])
        return None

    def update_chat_history(self, new_messages):
        updated_history = self.chat_history + new_messages
        conn = sqlite3.connect('users.db')
        c = conn.cursor()
        c.execute('UPDATE users SET previous_chat_history = ? WHERE id = ?',
                 (str(updated_history), self.id))
        conn.commit()
        conn.close()
        self.chat_history = updated_history  # Update in-memory

    def update_products_bought(self, new_products):
        updated_products = self.products_bought + new_products
        conn = sqlite3.connect('users.db')
        c = conn.cursor()
        c.execute('UPDATE users SET previous_products_bought = ? WHERE id = ?',
                 (str(updated_products), self.id))
        conn.commit()
        conn.close()
        self.products_bought = updated_products  # Update in-memory

# --- AI 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"],
)

embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

@st.cache_resource(show_spinner=False)
def load_data():
    products = Product.get_all()
    docs = []
    for p in products:
        content = f"Name: {p[1]}\nCategory: {p[2]}\nPrice: {p[3]}\nDescription: {p[4]}\nFeatures: {p[5]}\nStock: {p[6]}"
        metadata = {"id": p[0], "name": p[1], "category": p[2], "price": p[3], "stock": p[6]}
        docs.append(Document(page_content=content, metadata=metadata))
    
    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 get_current_inventory():
    """
    Get the current inventory with accurate stock levels.
    Returns a list of products with their details.
    """
    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    c.execute('SELECT id, name, category, price, description, features, stock FROM products WHERE stock > 0')
    in_stock_products = c.fetchall()
    conn.close()
    return in_stock_products

def verify_product_exists(product_name: str) -> dict:
    """
    Verify if a product exists in inventory and has stock.
    Returns a dict with product info or error message.
    """
    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    # Use flexible matching to find the product
    c.execute('SELECT id, name, category, price, stock FROM products WHERE name LIKE ?', 
             (f'%{product_name}%',))
    products = c.fetchall()
    conn.close()
    
    if not products:
        return {"exists": False, "message": f"Product '{product_name}' not found in inventory."}
    
    # Find best match
    best_match = None
    for p in products:
        if p[1].lower() == product_name.lower():
            best_match = p
            break
    
    # If no exact match, use the first similar match
    if not best_match and products:
        best_match = products[0]
    
    if best_match[4] <= 0:
        return {
            "exists": True, 
            "in_stock": False,
            "name": best_match[1],
            "message": f"Product '{best_match[1]}' exists but is out of stock."
        }
    
    return {
        "exists": True,
        "in_stock": True,
        "id": best_match[0],
        "name": best_match[1],
        "category": best_match[2],
        "price": best_match[3],
        "stock": best_match[4],
        "message": f"Product '{best_match[1]}' is in stock ({best_match[4]} available)."
    }

def retrieve_query(query: str):
    """
    Enhanced retrieval function that properly filters out products with zero stock.
    Returns only in-stock products unless explicitly querying for stock information.
    """
    # Get all relevant documents from the retriever
    retriever = load_data()
    docs = retriever.get_relevant_documents(query)
    
    # Store all retrieved docs for debug/admin purposes
    st.session_state.last_retrieved_all_docs = docs
    
    # Filter out products with zero stock unless explicitly checking stock
    if not any(term in query.lower() for term in ["stock", "inventory", "available"]):
        docs = [doc for doc in docs if doc.metadata.get("stock", 0) > 0]
    
    # Store the filtered docs that will be used by the agent
    st.session_state.last_retrieved_docs = docs
    
    # Add explicit message when no products found
    if not docs:
        return [{
            "page_content": f"No products matching '{query}' are currently in stock.",
            "metadata": {"notice": "no_products_found"}
        }]
    
    return docs

# Create tools for the agent
product_tool = Tool(
    name="product_retriever",
    func=retrieve_query,
    description="Useful for retrieving product information. Only returns in-stock products unless explicitly checking inventory status."
)

inventory_check_tool = StructuredTool.from_function(
    func=verify_product_exists,
    name="inventory_verifier",
    description="ALWAYS use this tool to verify if a specific product exists and has stock before recommending it",
    return_direct=False
)

def verify_response_products(response, current_inventory):
    """
    Verify that all products mentioned in the response exist in inventory.
    Add warnings if products are mentioned that don't exist or are out of stock.
    """
    inventory_names = [p[1].lower() for p in current_inventory]
    
    # Simple pattern to find potential product mentions (capitalized words)
    product_pattern = r'\b[A-Z][a-zA-Z0-9]+ (?:[A-Z][a-zA-Z0-9]+ )*(?:Laptop|Phone|Computer|Device|Tablet|Watch|Camera|Headphones|Speaker|Monitor|TV|Console)\b'
    potential_products = re.findall(product_pattern, response)
    
    warnings = []
    for product in potential_products:
        if product.lower() not in [name.lower() for name in inventory_names]:
            # Check if this is a real product that's out of stock or completely made up
            result = verify_product_exists(product)
            if result["exists"] and not result["in_stock"]:
                warnings.append(f"Warning: '{product}' was mentioned but is out of stock.")
            elif not result["exists"]:
                warnings.append(f"Warning: '{product}' was mentioned but doesn't exist in our inventory.")
    
    if warnings:
        st.warning("Possible inventory discrepancies detected in the AI response:")
        for warning in warnings:
            st.warning(warning)
            
    return response

# --- Admin Dashboard ---
def admin_dashboard():
    st.header("Admin Dashboard")
    
    with st.expander("Company Settings"):
        current_settings = Company.get_settings()
        with st.form("Company Settings Form"):
            name = st.text_input("Company Name", value=current_settings[1])
            business = st.text_input("Business", value=current_settings[2])
            agent_name = st.text_input("Agent Name", value=current_settings[3])
            key_features = st.text_area("Key Features", value=current_settings[4])
            
            if st.form_submit_button("Update Settings"):
                Company.update_settings(name, business, agent_name, key_features)
                st.success("Settings updated!")
                st.rerun()
    
    with st.expander("Product Management"):
        with st.form("Add Product"):
            st.subheader("Add New Product")
            name = st.text_input("Product Name")
            category = st.text_input("Category")
            price = st.number_input("Price", min_value=0.0)
            description = st.text_area("Description")
            features = st.text_area("Features")
            stock = st.number_input("Initial Stock", min_value=0, value=0)
            
            if st.form_submit_button("Add Product"):
                Product.add(name, category, price, description, features, stock)
                st.success("Product added!")
                load_data.clear()
                st.rerun()
        
        st.subheader("Manage Inventory")
        products = Product.get_all()
        if products:
            for p in products:
                cols = st.columns([3,2,2,1])
                cols[0].write(f"**{p[1]}** ({p[2]})")
                cols[1].write(f"Price: ${p[3]}")
                new_stock = cols[2].number_input(
                    "Stock", 
                    min_value=0, 
                    value=p[6],
                    key=f"stock_{p[0]}"
                )
                if new_stock != p[6]:
                    Product.update_stock(p[0], new_stock)
                    load_data.clear()
                    st.rerun()
                if cols[3].button("❌", key=f"del_{p[0]}"):
                    Product.delete(p[0])
                    load_data.clear()
                    st.rerun()
        else:
            st.info("No products found in database")

# --- Main App ---
def main():
    company_settings = Company.get_settings()
    company_name = company_settings[1]

    st.title("AI Sales Assistant 🤖")
    
    if 'user' not in st.session_state:
        st.session_state.user = None
        st.session_state.chat_history = []
        st.session_state.admin_mode = False
        st.session_state.last_retrieved_docs = []
        st.session_state.last_retrieved_all_docs = []
    
    # Authentication
    if not st.session_state.user and not st.session_state.admin_mode:
        st.header("Login/Register Admin")
        tab1, tab2, tab3 = st.tabs(["Login", "Register", "Admin"])
        
        with tab1:
            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):
                        st.session_state.user = user
                        st.session_state.chat_history = user.chat_history
                        st.rerun()
                    else:
                        st.error("Invalid credentials")
        
        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("Register"):
                    if User.get_by_username(new_user):
                        st.error("Username already exists")
                    else:
                        user = User.create(new_user, new_pass)
                        st.session_state.user = user
                        st.session_state.chat_history = []
                        st.rerun()
        
        with tab3:
            with st.form("Admin Login"):
                admin_pin = st.text_input("Admin PIN", type="password")
                if st.form_submit_button("Admin Login"):
                    if admin_pin == st.secrets["ADMIN_PIN"]:
                        st.session_state.admin_mode = True
                        st.rerun()
                    else:
                        st.error("Invalid Admin PIN")

    elif st.session_state.admin_mode:
        admin_dashboard()
        if st.button("Exit Admin Mode"):
            st.session_state.admin_mode = False
            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"])

        # Enhanced System Prompt
        company_settings = Company.get_settings()
        in_stock_products = Product.get_in_stock()
        current_products = "\n".join([f"- {p[1]} (${p[3]}, Stock: {p[6]})" for p in in_stock_products]) if in_stock_products else "No products currently in stock."
        
        system_prompt = f"""
        You are {company_settings[3]}, the AI Sales Assistant for {company_settings[1]} ({company_settings[2]}). Your primary role is to assist customers with product inquiries, make appropriate recommendations, and facilitate purchases while STRICTLY adhering to company policies.

        ## CRITICAL INVENTORY RULES
        - You may ONLY recommend products that exist in the current inventory list
        - ALWAYS check available stock BEFORE suggesting any product
        - DO NOT mention or suggest products with 0 stock
        - If a product exists but has 0 stock, inform the customer it's "currently out of stock"
        - If a product doesn't exist in our inventory, NEVER make up information about it

        ## Current Inventory Status
        {current_products}

        ## Company Profile
        - Company Name: {company_settings[1]}
        - Business Type: {company_settings[2]}
        - Key Features: {company_settings[4]}
        - Agent Name: {company_settings[3]}

        ## Inventory Management Policy
        1. **Stock Verification**:
           - BEFORE every product recommendation, verify the item exists in our current inventory
           - NEVER suggest out-of-stock items (stock = 0)
           - For low stock items (stock ≤ 3), explicitly mention: "Only [X] left in stock!"
           - If current_products is empty, inform the user: "I apologize, but our inventory system shows we currently have no products in stock. Please check back later."

        2. **Product Recommendations**:
           - ONLY recommend products from the current inventory list above
           - ALWAYS verify product existence and stock before recommendation
           - If asked about a specific product not in our inventory, respond with: "I apologize, we don't currently carry that item. As a {company_settings[2]}, we specialize in [relevant categories from our inventory]. May I suggest [alternative from current inventory with stock > 0]?"
           - When suggesting alternatives, ONLY suggest items with available stock

        3. **Purchase Process**:
           - Double-check product availability before processing any purchase
           - Generate payment link only after confirming stock > 0
           - Update inventory immediately after successful purchase
           - If stock becomes 0 during transaction, inform customer and suggest alternatives

        ## Conversation Flow Examples

        ### 1. Greeting & Need Assessment
        User: "Hi, I need a new laptop"
        You: "Hello! I'd be happy to help you find the perfect laptop from our current inventory. Could you tell me what you'll primarily be using it for and your budget range?"

        ### 2. Product Recommendation (In-Stock)
        User: "I need a gaming laptop under $1500"
        You: "Based on our current inventory, we have the XYZ Gaming Laptop available for $1399 (3 in stock). It features [key specs]. Would you like more details?"

        ### 3. Product Recommendation (Out-of-Stock)
        User: "Do you have ABC Smartphone?"  
        You: "I've checked our inventory, and I apologize, but we don't currently carry the ABC Smartphone. However, we do have the DEF Smartphone available for $599 (5 in stock), which offers similar features. Would you like me to tell you more about it?"

        ### 4. Out-of-Category Request
        User: "Do you sell kitchen appliances?"
        You: "I apologize, as a {company_settings[2]}, we specialize in [categories from current inventory]. We don't carry kitchen appliances. Here are some products from our current inventory that might interest you: [list 2-3 popular in-stock items]"

        ### 5. Purchase Process
        User: "I want to buy the XYZ Laptop"
        You: "Great choice! I've checked our inventory and the XYZ Laptop is available for $1399 (2 in stock). I can generate a secure payment link for you. Would you like to proceed with the purchase? [https://www.example.com/payment]"

        ## Communication Guidelines
        - Tone: Professional yet friendly (like a knowledgeable salesperson)
        - Language: Clear, concise, avoid technical jargon unless requested
        - Emojis: Use sparingly (1-2 per message max)
        - Branding: Consistently reference {company_settings[1]} when appropriate

        ## FINAL VERIFICATION CHECKLIST
        Before sending ANY response that recommends a product:
        1. Verify the product exists in current_products
        2. Confirm stock > 0
        3. Include accurate price and stock information
        4. Only if current_products is empty, inform user our inventory is currently empty

        Remember: Your primary responsibility is inventory accuracy. NEVER recommend products not in our current inventory or with 0 stock.
        """

        prompt = ChatPromptTemplate.from_messages([
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad")
        ])

        tools = [product_tool, inventory_check_tool]
        agent = create_tool_calling_agent(llm, tools, prompt)
        agent_executor = AgentExecutor(
            agent=agent, 
            tools=tools, 
            verbose=True,
            handle_parsing_errors=True
        )

        if prompt_input := st.chat_input("Type your message here..."):
            with st.chat_message("user"):
                st.write(prompt_input)

            with st.chat_message("assistant"):
                # Get current inventory for verification
                current_inventory = get_current_inventory()
                
                response = agent_executor.invoke({
                    "input": prompt_input,
                    "chat_history": [HumanMessage(content=msg["content"]) if msg["type"] == "human" else AIMessage(content=msg["content"]) 
                                   for msg in st.session_state.chat_history]
                })["output"]
                
                # Verify response integrity
                verified_response = verify_response_products(response, current_inventory)
                
                st.write(verified_response)
                
                # Enhanced purchase handling
                if "https://www.example.com/payment" in response:
                    # Extract product name from the agent's response
                    product_name = None
                    if "proceed with the purchase of" in response:
                        start = response.find("proceed with the purchase of") + len("proceed with the purchase of")
                        end = response.find("?", start)
                        product_name = response[start:end].strip()
                    
                    if product_name:
                        # Final check to confirm stock availability
                        product_check = verify_product_exists(product_name)
                        if not product_check["exists"]:
                            st.error(f"Error: Product '{product_name}' not found in our inventory.")
                        elif not product_check["in_stock"]:
                            st.error(f"Error: Product '{product_name}' is now out of stock.")
                        else:
                            # Process purchase with verified product information
                            product_id = product_check["id"]
                            try:
                                conn = sqlite3.connect('users.db')
                                c = conn.cursor()
                                
                                # Update product stock
                                c.execute('UPDATE products SET stock = stock - 1 WHERE id = ?', (product_id,))
                                
                                # Update user's purchase history
                                c.execute('SELECT previous_products_bought FROM users WHERE id = ?', 
                                        (st.session_state.user.id,))
                                current_purchases = c.fetchone()[0]
                                
                                updated_purchases = eval(current_purchases) + [product_name] if current_purchases else [product_name]
                                
                                c.execute('UPDATE users SET previous_products_bought = ? WHERE id = ?',
                                        (str(updated_purchases), st.session_state.user.id))
                                
                                conn.commit()
                                st.session_state.user.products_bought.append(product_name)
                                st.success(f"Successfully purchased {product_name}! Stock updated.")
                                load_data.clear()  # Refresh inventory data
                                
                                # Check for low stock
                                c.execute('SELECT stock FROM products WHERE id = ?', (product_id,))
                                remaining_stock = c.fetchone()[0]
                                if remaining_stock <= 3 and remaining_stock > 0:
                                    st.warning(f"Low stock alert: Only {remaining_stock} units of {product_name} remaining!")
                            except Exception as e:
                                st.error(f"Error processing purchase: {str(e)}")
                                if conn:
                                    conn.rollback()
                            finally:
                                if conn:
                                    conn.close()

            new_messages = [
                {"type": "human", "content": prompt_input},
                {"type": "ai", "content": verified_response}
            ]
            st.session_state.user.update_chat_history(new_messages)
            st.session_state.chat_history += new_messages

if __name__ == "__main__":
    main()