Update app.py

app.py

@@ -1,951 +1,518 @@
+
 import os
 import sqlite3
 import streamlit as st
 from werkzeug.security import generate_password_hash, check_password_hash
-import logging # Import logging
 
-# Langchain specific imports
 from langchain_groq import ChatGroq
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.document_loaders.csv_loader import CSVLoader # Although not used directly for DB, keep if needed later
+from langchain_community.document_loaders.csv_loader import CSVLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import FAISS # Using FAISS for persistence/efficiency
-# Or use InMemoryVectorStore if preferred and persistence isn't critical:
-# from langchain_core.vectorstores import InMemoryVectorStore
-from langchain_core.vectorstores import VectorStore # Base class for type hinting
+from langchain_core.vectorstores import InMemoryVectorStore
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.tools import Tool
 from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain_core.messages import HumanMessage, AIMessage
 from langchain.docstore.document import Document
 
-# --- Logging Setup ---
-# Configure logging for better debugging, especially for agent actions
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
 # --- Database Setup ---
-# Use a constant for the database file name
-DB_FILE = 'users_and_products.db'
-
-@st.cache_resource(show_spinner="Initializing Database...")
+@st.cache_resource
 def init_db():
-    """Initializes the SQLite database and returns a connection."""
-    try:
-        conn = sqlite3.connect(DB_FILE, 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 DEFAULT 1, -- Ensure only one row
-                      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(stock >= 0))''') # Ensure stock >= 0
-
-        # Check and update products schema if 'stock' column is missing (for backward compatibility)
-        c.execute("PRAGMA table_info(products)")
-        columns = [column[1] for column in c.fetchall()]
-        if 'stock' not in columns:
-            logger.warning("Adding missing 'stock' column to products table.")
-            c.execute('ALTER TABLE products ADD COLUMN stock INTEGER NOT NULL DEFAULT 0 CHECK(stock >= 0)')
-
-        # Insert default company settings if empty
-        c.execute('SELECT COUNT(*) FROM company_settings')
-        if c.fetchone()[0] == 0:
-            logger.info("Inserting default company settings.")
-            c.execute('''INSERT OR IGNORE INTO company_settings
-                         (id, name, business, agent_name, key_features)
-                         VALUES (?, ?, ?, ?, ?)''',
-                      (1, # Explicitly set ID to 1
-                       'TechElectronics',
-                       'Consumer Electronics Retailer',
-                       'Alex',
-                       'Cutting-edge technology, Competitive pricing, Excellent customer service'))
-
-        conn.commit()
-        logger.info("Database initialized successfully.")
-        return conn
-    except sqlite3.Error as e:
-        logger.error(f"Database initialization failed: {e}")
-        st.error(f"Fatal Error: Could not initialize the database. Please check logs. Error: {e}")
-        st.stop() # Stop execution if DB fails
-    # No need to return conn here as it's handled by cache_resource
-
-# Get a connection (will be cached after first call)
-# Note: It's generally better practice to open/close connections per transaction/request
-# in web apps, but for Streamlit's model, caching the connection object is common,
-# though requires `check_same_thread=False`. Be mindful of potential concurrency issues.
-# A safer approach might involve a connection pool or context managers within functions.
-# For simplicity in this example, we'll use the cached connection but wrap DB ops in try/finally.
-
-def get_db_conn():
-    """Gets a connection to the database."""
-    try:
-        return sqlite3.connect(DB_FILE, check_same_thread=False)
-    except sqlite3.Error as e:
-        logger.error(f"Failed to connect to database: {e}")
-        st.error(f"Error connecting to the database: {e}")
-        return None
+    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():
-        """Retrieves company settings from the database."""
-        conn = get_db_conn()
-        if not conn: return None
-        try:
-            c = conn.cursor()
-            c.execute('SELECT name, business, agent_name, key_features FROM company_settings WHERE id = 1')
-            settings = c.fetchone()
-            # Return a default tuple if no settings found, preventing downstream errors
-            return settings if settings else ('Default Company', 'Default Business', 'AI Assistant', 'Default Features')
-        except sqlite3.Error as e:
-            logger.error(f"Error getting company settings: {e}")
-            return None # Indicate error
-        finally:
-            if conn: conn.close()
+        c = conn.cursor()
+        c.execute('SELECT * FROM company_settings LIMIT 1')
+        return c.fetchone()
 
     @staticmethod
     def update_settings(name, business, agent_name, key_features):
-        """Updates company settings in the database."""
-        conn = get_db_conn()
-        if not conn: return False
-        try:
-            c = conn.cursor()
-            # Use INSERT OR REPLACE to handle the case where the row might not exist yet (though init_db should prevent this)
-            c.execute('''INSERT OR REPLACE INTO company_settings
-                         (id, name, business, agent_name, key_features)
-                         VALUES (1, ?, ?, ?, ?)''',
-                      (name, business, agent_name, key_features))
-            conn.commit()
-            logger.info("Company settings updated.")
-            return True
-        except sqlite3.Error as e:
-            logger.error(f"Error updating company settings: {e}")
-            conn.rollback()
-            return False
-        finally:
-            if conn: conn.close()
+        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():
-        """Retrieves all products from the database."""
-        conn = get_db_conn()
-        if not conn: return []
-        try:
-            c = conn.cursor()
-            c.execute('SELECT id, name, category, price, description, features, stock FROM products ORDER BY name')
-            return c.fetchall()
-        except sqlite3.Error as e:
-            logger.error(f"Error getting all products: {e}")
-            return [] # Return empty list on error
-        finally:
-            if conn: conn.close()
+        c = conn.cursor()
+        c.execute('SELECT * FROM products')
+        return c.fetchall()
 
     @staticmethod
     def get_by_id(product_id):
-        """Retrieves a single product by its ID."""
-        conn = get_db_conn()
-        if not conn: return None
-        try:
-            c = conn.cursor()
-            c.execute('SELECT id, name, category, price, description, features, stock FROM products WHERE id = ?', (product_id,))
-            return c.fetchone()
-        except sqlite3.Error as e:
-            logger.error(f"Error getting product by ID {product_id}: {e}")
-            return None
-        finally:
-            if conn: conn.close()
+        c = conn.cursor()
+        c.execute('SELECT * FROM products WHERE id = ?', (product_id,))
+        return c.fetchone()
 
     @staticmethod
     def add(name, category, price, description, features, stock):
-        """Adds a new product to the database."""
-        if not all([name, category, description, features]) or price < 0 or stock < 0:
-             st.error("Invalid product data. Please check all fields.")
-             return False
-        conn = get_db_conn()
-        if not conn: return False
-        try:
-            c = conn.cursor()
-            c.execute('''INSERT INTO products
-                         (name, category, price, description, features, stock)
-                         VALUES (?, ?, ?, ?, ?, ?)''',
-                      (name, category, price, description, features, stock))
-            conn.commit()
-            logger.info(f"Product '{name}' added successfully.")
-            # Clear cache after adding product
-            load_vector_store.clear()
-            return True
-        except sqlite3.IntegrityError:
-             logger.error(f"Product '{name}' might already exist or violates constraints.")
-             st.error(f"Product '{name}' might already exist.")
-             conn.rollback()
-             return False
-        except sqlite3.Error as e:
-            logger.error(f"Error adding product '{name}': {e}")
-            st.error(f"Database error adding product: {e}")
-            conn.rollback()
-            return False
-        finally:
-            if conn: conn.close()
+        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):
-        """Deletes a product from the database by ID."""
-        conn = get_db_conn()
-        if not conn: return False
-        try:
-            c = conn.cursor()
-            c.execute('DELETE FROM products WHERE id=?', (product_id,))
-            conn.commit()
-            logger.info(f"Product ID {product_id} deleted.")
-            # Clear cache after deleting product
-            load_vector_store.clear()
-            return True
-        except sqlite3.Error as e:
-            logger.error(f"Error deleting product ID {product_id}: {e}")
-            conn.rollback()
-            return False
-        finally:
-            if conn: conn.close()
+        c = conn.cursor()
+        c.execute('DELETE FROM products WHERE id=?', (product_id,))
+        conn.commit()
 
     @staticmethod
     def update_stock(product_id, new_stock):
-        """Updates the stock level for a specific product."""
-        if new_stock < 0:
-            logger.warning(f"Attempted to set negative stock ({new_stock}) for product ID {product_id}.")
-            st.warning("Stock cannot be negative.")
-            return False
-        conn = get_db_conn()
-        if not conn: return False
-        try:
-            c = conn.cursor()
-            c.execute('UPDATE products SET stock=? WHERE id=?', (new_stock, product_id))
-            conn.commit()
-            logger.debug(f"Stock for product ID {product_id} updated to {new_stock}.")
-             # Clear cache after updating stock
-            load_vector_store.clear()
-            return True
-        except sqlite3.Error as e:
-            logger.error(f"Error updating stock for product ID {product_id}: {e}")
-            conn.rollback()
-            return False
-        finally:
-            if conn: conn.close()
+        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_hash, chat_history_str=None, products_bought_str=None):
+    def __init__(self, id, username, password, chat_history=None, products_bought=None):
         self.id = id
         self.username = username
-        self.password_hash = password_hash # Store the hash directly
-
-        # Safely evaluate stored string representations of lists
-        try:
-            self.chat_history = eval(chat_history_str) if chat_history_str else []
-            if not isinstance(self.chat_history, list): self.chat_history = []
-        except Exception as e:
-            logger.warning(f"Error parsing chat history for user {username}: {e}. Resetting to empty list.")
-            self.chat_history = []
-
-        try:
-            self.products_bought = eval(products_bought_str) if products_bought_str else []
-            if not isinstance(self.products_bought, list): self.products_bought = []
-        except Exception as e:
-            logger.warning(f"Error parsing products bought for user {username}: {e}. Resetting to empty list.")
-            self.products_bought = []
-
+        self.password = password
+        self.chat_history = chat_history or []
+        self.products_bought = products_bought or []
 
     @classmethod
     def create(cls, username, password):
-        """Creates a new user, hashes the password, and saves to DB."""
-        if not username or not password:
-             st.error("Username and password cannot be empty.")
-             return None
         hashed_pw = generate_password_hash(password)
-        conn = get_db_conn()
-        if not conn: return None
-        try:
-            c = conn.cursor()
-            c.execute('INSERT INTO users (username, password, previous_chat_history, previous_products_bought) VALUES (?, ?, ?, ?)',
-                      (username, hashed_pw, '[]', '[]')) # Initialize with empty lists as strings
-            user_id = c.lastrowid
-            conn.commit()
-            logger.info(f"User '{username}' created successfully with ID {user_id}.")
-            return cls(user_id, username, hashed_pw) # Return new User instance
-        except sqlite3.IntegrityError:
-            logger.warning(f"Attempted to create user with existing username: {username}")
-            st.error("Username already exists. Please choose a different one.")
-            conn.rollback()
-            return None
-        except sqlite3.Error as e:
-            logger.error(f"Database error creating user '{username}': {e}")
-            st.error(f"Database error during registration: {e}")
-            conn.rollback()
-            return None
-        finally:
-            if conn: conn.close()
+        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):
-        """Retrieves a user by username from the database."""
-        conn = get_db_conn()
-        if not conn: return None
-        try:
-            c = conn.cursor()
-            c.execute('SELECT id, username, password, previous_chat_history, previous_products_bought FROM users WHERE username = ?', (username,))
-            user_data = c.fetchone()
-            if user_data:
-                # Pass the raw string data to the constructor for parsing
-                return cls(user_data[0], user_data[1], user_data[2], user_data[3], user_data[4])
-            else:
-                return None # User not found
-        except sqlite3.Error as e:
-            logger.error(f"Database error retrieving user '{username}': {e}")
-            return None
-        finally:
-            if conn: conn.close()
-
-    def _update_field(self, field_name, updated_value_list):
-        """Helper method to update a list field (chat history or products) in the DB."""
-        conn = get_db_conn()
-        if not conn: return False
-        try:
-            c = conn.cursor()
-            # Store the list as its string representation
-            c.execute(f'UPDATE users SET {field_name} = ? WHERE id = ?',
-                      (str(updated_value_list), self.id))
-            conn.commit()
-            logger.debug(f"User {self.id}'s {field_name} updated in DB.")
-            return True
-        except sqlite3.Error as e:
-            logger.error(f"Database error updating {field_name} for user {self.id}: {e}")
-            conn.rollback()
-            return False
-        finally:
-            if conn: conn.close()
+        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):
-        """Appends new messages to the user's chat history and updates the DB."""
-        # Ensure new_messages is a list
-        if not isinstance(new_messages, list):
-             logger.error("Invalid format for new_messages. Expected a list.")
-             return
-
-        # Append new messages to the existing list
         updated_history = self.chat_history + new_messages
-        if self._update_field('previous_chat_history', updated_history):
-            self.chat_history = updated_history # Update in-memory object only on successful DB update
-
-    def update_products_bought(self, new_product_name):
-        """Adds a new product name to the user's bought list and updates the DB."""
-         # Ensure new_product_name is a string
-        if not isinstance(new_product_name, str):
-             logger.error("Invalid format for new_product_name. Expected a string.")
-             return
+        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
 
-        updated_products = self.products_bought + [new_product_name] # Append the new product name
-        if self._update_field('previous_products_bought', updated_products):
-            self.products_bought = updated_products # Update in-memory object
+    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 ---
-# Load API Key from secrets
-try:
-    GROQ_API_KEY = st.secrets["GROQ_API_KEY"]
-    ADMIN_PIN = st.secrets["ADMIN_PIN"]
-    os.environ["GROQ_API_KEY"] = GROQ_API_KEY
-except KeyError as e:
-    st.error(f"Missing secret: {e}. Please configure GROQ_API_KEY and ADMIN_PIN in .streamlit/secrets.toml")
-    st.stop()
-
-# Initialize LLM
-# Consider adding error handling for LLM initialization
-try:
-    llm = ChatGroq(
-        temperature=0.1, # Low temperature for more factual responses
-        model_name="llama3-8b-8192", # Or other suitable model
-        # api_key=GROQ_API_KEY, # Handled by environment variable
-    )
-except Exception as e:
-    st.error(f"Failed to initialize the Language Model: {e}")
-    logger.error(f"LLM Initialization failed: {e}")
-    st.stop()
-
-
-# Initialize Embeddings Model
-# Using a smaller, faster model suitable for in-memory vector stores
-try:
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",
-                                       model_kwargs={'device': 'cpu'}) # Explicitly use CPU if GPU issues arise
-except Exception as e:
-    st.error(f"Failed to initialize the Embeddings Model: {e}")
-    logger.error(f"Embeddings Initialization failed: {e}")
-    st.stop()
-
-# --- Vector Store and Retriever ---
-# Path for FAISS index persistence
-FAISS_INDEX_PATH = "faiss_product_index"
-
-@st.cache_resource(show_spinner="Loading Product Knowledge...")
-def load_vector_store() -> VectorStore:
-    """Loads products from DB, creates documents, splits them, and builds/loads a FAISS vector store."""
-    logger.info("Attempting to load or build vector store.")
-    products = Product.get_all()
-    if not products:
-        logger.warning("No products found in database to build vector store.")
-        # Return an empty store or handle appropriately
-        # For simplicity, creating an empty one here, but might need better handling
-        return FAISS.from_texts(["No products available"], embeddings) # Create dummy store
+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:
-        # Create detailed content string for each product
-        content = (f"Product Name: {p[1]}\n"
-                   f"Category: {p[2]}\n"
-                   f"Price: ${p[3]:.2f}\n" # Format price
-                   f"Stock: {p[6]} units available\n" # Crucial: Include current stock
-                   f"Description: {p[4]}\n"
-                   f"Features: {p[5]}")
-        # Metadata includes ID, name, category, price, and stock for potential filtering/retrieval checks
+        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()
 
-    if not docs:
-         logger.warning("No documents created from products.")
-         return FAISS.from_texts(["No products available"], embeddings)
+retriever = load_data()
 
-    # Split documents into smaller chunks
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) # Adjusted chunk size
-    splits = text_splitter.split_documents(docs)
+def retrieve_query(query: str):
+    docs = retriever.get_relevant_documents(query)
+    st.session_state.last_retrieved_docs = docs
+    return docs
 
-    # Create or load FAISS index
-    # Note: FAISS persistence is basic. For production, consider more robust vector DBs.
-    if os.path.exists(FAISS_INDEX_PATH):
-        try:
-            logger.info(f"Loading existing FAISS index from {FAISS_INDEX_PATH}")
-            vectorstore = FAISS.load_local(FAISS_INDEX_PATH, embeddings, allow_dangerous_deserialization=True) # Required for FAISS loading
-            # Optional: Update the index if products changed significantly?
-            # vectorstore.add_documents(splits) # Could lead to duplicates if not managed carefully
-            logger.info("FAISS index loaded.")
-        except Exception as e:
-            logger.error(f"Error loading FAISS index: {e}. Rebuilding index.")
-            vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
-            vectorstore.save_local(FAISS_INDEX_PATH)
-            logger.info("New FAISS index built and saved.")
-    else:
-        logger.info("Building new FAISS index.")
-        vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
-        vectorstore.save_local(FAISS_INDEX_PATH)
-        logger.info(f"New FAISS index built and saved to {FAISS_INDEX_PATH}.")
-
-    return vectorstore
-
-# Load the vector store (cached)
-vector_store = load_vector_store()
-retriever = vector_store.as_retriever(search_kwargs={"k": 5}) # Retrieve top 5 relevant chunks
-
-# --- Tool Definition ---
-def retrieve_query(query: str) -> list[Document]:
-    """Retrieves relevant product documents based on the query using the vector store."""
-    logger.info(f"Retrieving documents for query: '{query}'")
-    try:
-        # Use the globally loaded retriever
-        docs = retriever.invoke(query)
-        # Log retrieved docs for debugging (optional)
-        # logger.debug(f"Retrieved {len(docs)} documents: {[doc.metadata for doc in docs]}")
-        # Store last retrieved docs in session state IF NEEDED for purchase logic (alternative to re-querying)
-        # st.session_state.last_retrieved_docs = docs
-        return docs
-    except Exception as e:
-        logger.error(f"Error during document retrieval for query '{query}': {e}")
-        return [] # Return empty list on error
-
-# Create the Langchain Tool
-product_retriever_tool = Tool(
+tool = Tool(
     name="product_retriever",
     func=retrieve_query,
-    description="Mandatory tool to get CURRENT product information. Use it to find products by name or category, check features, price, and MOST IMPORTANTLY, current stock levels BEFORE recommending or confirming anything to the customer."
+    description="Useful for retrieving product information including current stock levels"
 )
 
 # --- Admin Dashboard ---
 def admin_dashboard():
-    """Renders the admin dashboard for managing settings and products."""
     st.header("Admin Dashboard")
-
-    # Logout Button
-    if st.sidebar.button("Exit Admin Mode"):
-        st.session_state.admin_mode = False
-        st.rerun()
-
-    # Company Settings Management
-    with st.expander("Company Settings", expanded=True):
+    
+    with st.expander("Company Settings"):
         current_settings = Company.get_settings()
-        if current_settings is None:
-            st.error("Could not load company settings.")
-            return # Stop rendering if settings failed to load
-
         with st.form("Company Settings Form"):
-            # Use values from fetched settings or defaults if None
-            name = st.text_input("Company Name", value=current_settings[0])
-            business = st.text_input("Business", value=current_settings[1])
-            agent_name = st.text_input("Agent Name", value=current_settings[2])
-            key_features = st.text_area("Key Features", value=current_settings[3])
-
-            submitted = st.form_submit_button("Update Settings")
-            if submitted:
-                if Company.update_settings(name, business, agent_name, key_features):
-                    st.success("Settings updated successfully!")
-                    # No need to clear cache here unless settings affect prompts drastically
-                else:
-                    st.error("Failed to update settings.")
-
-    # Product Management
-    with st.expander("Product Management", expanded=True):
-        # Add New Product Form
-        with st.form("Add Product 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!")
+    
+    with st.expander("Product Management"):
+        with st.form("Add Product"):
             st.subheader("Add New Product")
-            p_name = st.text_input("Product Name*")
-            p_category = st.text_input("Category*")
-            p_price = st.number_input("Price*", min_value=0.0, format="%.2f")
-            p_description = st.text_area("Description*")
-            p_features = st.text_area("Features*")
-            p_stock = st.number_input("Initial Stock*", min_value=0, value=0, step=1)
-
-            add_submitted = st.form_submit_button("Add Product")
-            if add_submitted:
-                if not all([p_name, p_category, p_description, p_features]):
-                     st.error("Please fill in all required fields marked with *.")
-                elif Product.add(p_name, p_category, p_price, p_description, p_features, p_stock):
-                    st.success(f"Product '{p_name}' added!")
-                    # load_vector_store.clear() # Cache is cleared inside Product.add
-                    st.rerun() # Rerun to refresh the product list below
-                else:
-                    # Error message displayed within Product.add
-                    pass
-
-        # Manage Existing Products / Inventory
+            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.subheader("Manage Inventory")
         products = Product.get_all()
-        if not products:
-            st.info("No products found in the database. Add products using the form above.")
-        else:
-            st.markdown("""---""")
+        if products:
             for p in products:
-                # Use columns for better layout
-                col1, col2, col3, col4 = st.columns([3, 2, 2, 1]) # Adjust ratios as needed
-                with col1:
-                    st.write(f"**{p[1]}**") # Name
-                    st.caption(f"Category: {p[2]}") # Category
-                with col2:
-                    st.write(f"Price: ${p[3]:.2f}") # Price
-                with col3:
-                    # Use a unique key for each number input
-                    new_stock = st.number_input(
-                        "Stock",
-                        min_value=0,
-                        value=p[6], # Current stock (index 6)
-                        key=f"stock_{p[0]}", # Unique key using product ID (index 0)
-                        step=1,
-                        label_visibility="collapsed" # Hide label for cleaner look
-                    )
-                    # Update stock if value changes
-                    if new_stock != p[6]:
-                        if Product.update_stock(p[0], new_stock):
-                            st.toast(f"Stock for '{p[1]}' updated to {new_stock}.")
-                            # load_vector_store.clear() # Cache cleared inside update_stock
-                            st.rerun() # Rerun to reflect change immediately
-                        else:
-                             st.toast(f"Failed to update stock for '{p[1]}'.", icon="❌")
-
-
-                with col4:
-                    # Use a unique key for each delete button
-                    if st.button("❌ Delete", key=f"del_{p[0]}", help=f"Delete {p[1]}"):
-                        if Product.delete(p[0]):
-                            st.success(f"Product '{p[1]}' deleted.")
-                            # load_vector_store.clear() # Cache cleared inside delete
-                            st.rerun() # Rerun to refresh list
-                        else:
-                            st.error(f"Failed to delete product '{p[1]}'.")
-                st.markdown("""---""") # Separator
+                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)
+                    st.rerun()
+                if cols[3].button("❌", key=f"del_{p[0]}"):
+                    Product.delete(p[0])
+                    st.rerun()
+        else:
+            st.info("No products found in database")
 
-# --- Main App Logic ---
+# --- Main App ---
 def main():
-    """Main function to run the Streamlit application."""
-    # Initialize Database (cached)
-    init_db()
+    company_settings = Company.get_settings()
+    company_name = company_settings[1]
 
-    # Initialize session state variables if they don't exist
+    st.title("AI Sales Assistant 🤖")
+    
     if 'user' not in st.session_state:
-        st.session_state.user = None # Stores the logged-in User object
-    if 'chat_history' not in st.session_state:
-         # This will be populated from the user object upon login
-         # It primarily serves as the display buffer for the current session UI
+        st.session_state.user = None
         st.session_state.chat_history = []
-    if 'admin_mode' not in st.session_state:
         st.session_state.admin_mode = False
-    # 'last_retrieved_docs' could be used but re-querying before purchase is safer
-    # if 'last_retrieved_docs' not in st.session_state:
-    #     st.session_state.last_retrieved_docs = []
-
-    # --- Page Title ---
-    # Fetch company name for title (handle potential None)
-    company_settings = Company.get_settings()
-    company_name = company_settings[0] if company_settings else "Sales Assistant"
-    st.title(f"{company_name} AI Sales Assistant 🤖")
-
-
-    # --- Authentication / Mode Handling ---
-    if st.session_state.admin_mode:
-        admin_dashboard() # Show admin view
-    elif not st.session_state.user:
-        # Show Login/Register/Admin Access view
-        st.header("Login / Register / Admin Access")
-        login_tab, register_tab, admin_tab = st.tabs(["Login", "Register", "Admin Login"])
-
-        with login_tab:
-            with st.form("Login Form"):
+        st.session_state.last_retrieved_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")
-                login_submitted = st.form_submit_button("Login")
-                if login_submitted:
+                if st.form_submit_button("Login"):
                     user = User.get_by_username(username)
-                    if user and check_password_hash(user.password_hash, password):
+                    if user and check_password_hash(user.password, password):
                         st.session_state.user = user
-                        # Load user's history into session state for display
                         st.session_state.chat_history = user.chat_history
-                        logger.info(f"User '{username}' logged in.")
                         st.rerun()
                     else:
-                        st.error("Invalid username or password.")
-
-        with register_tab:
-            with st.form("Register Form"):
-                new_username = st.text_input("Choose Username")
-                new_password = st.text_input("Choose Password", type="password")
-                register_submitted = st.form_submit_button("Register")
-                if register_submitted:
-                    user = User.create(new_username, new_password)
-                    if user:
+                        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 = [] # Start with empty history for new user
-                        logger.info(f"User '{new_username}' registered and logged in.")
-                        st.success("Registration successful! You are now logged in.")
+                        st.session_state.chat_history = []
                         st.rerun()
-                    # Error messages handled within User.create
-
-        with admin_tab:
-            with st.form("Admin Login Form"):
-                admin_pin_input = st.text_input("Admin PIN", type="password")
-                admin_login_submitted = st.form_submit_button("Admin Login")
-                if admin_login_submitted:
-                    # Compare with the PIN loaded from secrets
-                    if admin_pin_input == ADMIN_PIN:
+        
+        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
-                        logger.info("Admin mode activated.")
                         st.rerun()
                     else:
-                        st.error("Invalid Admin PIN.")
-    else:
-        # --- Logged-in User Chat Interface ---
-        user: User = st.session_state.user # Get the logged-in user object
-        company_settings = Company.get_settings() # Fetch fresh settings
-        if not company_settings:
-            st.error("Critical Error: Could not load company settings for the chat.")
-            st.stop()
-
-        company_name = company_settings[0]
-        agent_name = company_settings[2]
-
-        st.sidebar.header(f"Welcome, {user.username}!")
-        if st.sidebar.button("Logout"):
-            logger.info(f"User '{user.username}' logged out.")
-            # Optionally save final chat history before logging out if needed
-            # user.update_chat_history([]) # Example: Save current session history (already done per message)
-            st.session_state.user = None
-            st.session_state.chat_history = []
+                        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()
 
-        st.sidebar.markdown("---")
-        st.sidebar.subheader("Purchase History")
-        if user.products_bought:
-            for item in user.products_bought:
-                st.sidebar.markdown(f"- {item}")
-        else:
-            st.sidebar.info("No purchases recorded yet.")
-
-
-        st.subheader(f"Chat with {agent_name}")
+    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 from session state
+        # Display 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"])
+            if msg["type"] == "human":
+                with st.chat_message("user"):
+                    st.write(msg["content"])
+            else:
+                with st.chat_message("assistant"):
+                    st.write(msg["content"])
 
-        # --- Define System Prompt (Using the improved version) ---
+        # Enhanced System Prompt
+        company_settings = Company.get_settings()
+        current_products = "\n".join([f"- {p[1]} (${p[3]}, Stock: {p[6]})" for p in Product.get_all()])
         system_prompt = f"""
-        **Your Identity:**
-        You are {company_settings[2]}, a friendly and highly professional AI Sales Assistant for {company_settings[0]}.
-        - **Company:** {company_settings[0]}
-        - **Business:** {company_settings[1]} (Consumer Electronics Retailer)
-        - **Key Strengths:** {company_settings[3]}
-
-        **Your Primary Goal:**
-        Assist customers effectively by understanding their needs, providing accurate product information based *only* on available inventory, guiding them through the sales process, and ensuring a positive customer experience.
-
-        **Core Operating Principles:**
-
-        1.  **Tool Reliance (Mandatory):**
-            - You MUST use the `product_retriever` tool to get **current and accurate** information about products, including name, description, features, price, and **especially stock levels**.
-            - **When to use the tool:**
-                - When asked about specific products or product categories.
-                - Before recommending ANY product to confirm its availability (stock > 0).
-                - Before confirming the price or stock level of a product.
-                - Before initiating the purchase process for a specific item.
-            - If the tool returns no relevant products for a query, state that you don't have information on that specific item or category currently available.
-            - If the tool fails unexpectedly, apologize and mention a temporary inability to fetch product details.
-
-        2.  **Inventory-Driven Recommendations:**
-            - **Only recommend products confirmed to be IN STOCK (stock > 0) via the `product_retriever` tool.**
-            - **Never** suggest, discuss, or confirm availability for products with stock <= 0.
-            - **Never** discuss or recommend products not part of {company_settings[0]}'s inventory (i.e., not returnable by the tool).
-
-        3.  **Stock Level Communication:**
-            - If a retrieved product has low stock (e.g., stock <= 3 BUT > 0), inform the customer using the "Low stock items" template.
-            - If a customer asks about a product that the tool confirms is out of stock (stock = 0), use the "Out-of-stock items" template and suggest relevant, *in-stock* alternatives retrieved via the tool.
-
-        4.  **User Context Awareness (Optional but Recommended):**
-            - You have access to the user's chat history. Use it to understand the context of the current conversation.
-            - The user object may contain `previous_products_bought`. If available ({user.products_bought}), you *can* use this information to tailor recommendations (e.g., accessories for past purchases), but **always prioritize current stock availability** confirmed via the tool.
-
-        **Sales Conversation Flow:**
-
-        1.  **Greeting & Needs Assessment:** Start with a warm, professional welcome. Ask open-ended questions to understand the customer's requirements, budget, and preferences. Actively listen.
-        2.  **Product Consultation & Recommendation:**
-            - Based on needs, query the `product_retriever` tool.
-            - Recommend **only suitable, in-stock** products.
-            - Clearly state product name, key features/benefits, and **tool-retrieved price and stock status**.
-            - If comparing products, use tool-retrieved data for both.
-        3.  **Handling Questions & Objections:** Address customer queries and concerns politely and factually, using information obtained from the `product_retriever` tool where applicable. If a requested product is unavailable or unsuitable, explain why and suggest appropriate, *in-stock* alternatives.
-        4.  **Purchase Process Guidance:**
-            - When a customer expresses clear intent to buy a *specific* product:
-                a. **Re-verify stock** using the `product_retriever` tool one last time. Crucially check the **stock** field in the retrieved document metadata.
-                b. If retrieved stock > 0, confirm the product name and **tool-retrieved price**.
-                c. Ask if they wish to proceed using the "For payment process" template.
-                d. If they agree, provide the pre-defined payment link: `https://www.example.com/payment`
-                e. **Crucially state:** "Please note: Stock is confirmed now, and inventory will be updated once you complete the process via the link."
-            - If retrieved stock is 0 upon re-verification, apologize and inform them it just went out of stock, suggesting alternatives if possible.
-
-        **Communication Style:**
-        - Tone: Professional, friendly, helpful, empathetic, patient.
-        - Language: Clear, concise, easy-to-understand. Avoid jargon unless explained.
-        - Branding: Consistently represent {company_settings[0]}.
-        - Formatting: Use bullet points or short paragraphs for readability.
-        - Emojis: Sparingly (max 1-2 per message) and appropriately. 😊
-
-        **Response Templates (Adapt as needed):**
-
-        * **Out-of-stock items:** "I apologize, but the [product name retrieved by tool] is currently out of stock (Tool shows 0 units). However, based on your needs, I found the [alternative in-stock product name retrieved by tool] which is available (Tool shows [stock] units) and has similar features like [mention 1-2 key features]. Would you like to know more about it?"
-        * **Low stock items:** "Just letting you know, we have limited availability for the [product name retrieved by tool]. The tool shows only [X retrieved stock] units left in stock. I'd recommend deciding soon if you're interested!"
-        * **Product not carried/found:** "I checked our inventory using the product tool, but I couldn't find [queried item] or anything similar in the [category if specified]. As a {company_settings[1]}, we specialize in [mention core categories]. Perhaps I can help you find something else from our available electronics?"
-        * **For payment process:** "Great choice! I've just double-checked with the tool, and the [product name retrieved by tool] is currently in stock. The price is $[price retrieved by tool]. Are you ready to proceed with the purchase? I can provide the secure payment link."
-        * **After providing payment link:** "Okay, here is the secure link to complete your purchase for the [product name]: https://www.example.com/payment. Please note: Stock is confirmed now, and inventory will be updated once you complete the process via the link."
-
-        **Key Constraints - Adhere Strictly:**
-        - **ALWAYS use the `product_retriever` tool for product details/stock BEFORE presenting information to the user.** Do not rely on memory or previous turns for stock levels. Check the 'stock' value in the retrieved document metadata.
-        - **Stick to the product inventory of {company_settings[0]}.** Do not discuss competitors or unavailable items unless suggesting in-stock alternatives.
-        - **Follow the purchase process precisely, especially the final stock check using the tool before offering the payment link.**
-        - **Prioritize accuracy and helpfulness.**
+        You are {company_settings[3]}, the AI Sales Assistant for {company_settings[1]} ({company_settings[2]}).
+
+        Company Profile:
+        - Company Name: {company_settings[1]}
+        - Business: {company_settings[2]}
+        - Key Features: {company_settings[4]}
+
+        Product Policy & Inventory Management:
+        1. Product Recommendations:
+           - Only recommend products with available stock (stock > 0)
+           - Always check current stock levels before making recommendations
+           - For out-of-stock items, suggest similar in-stock alternatives
+           - Never promote or discuss products not in our inventory
+
+        2. Stock Management:
+           - Track real-time inventory levels for all products
+           - Update stock automatically after successful purchases
+           - Alert customers about low stock items (stock ≤ 3)
+           - Inform when items are temporarily out of stock
+
+        Current Inventory:
+        {current_products}
+
+        Sales Process & Conversation Flow:
+        1. Greeting & Need Assessment
+           - Warm, professional welcome
+           - Ask relevant questions to understand customer needs
+           - Listen actively and acknowledge customer preferences
+
+        2. Product Consultation
+           - Recommend suitable products based on customer needs
+           - Highlight relevant features and benefits
+           - Provide accurate pricing and stock information
+           - Compare products when appropriate
+
+        3. Objection Handling
+           - Address concerns professionally
+           - Provide factual information
+           - Offer alternatives when necessary
+           - Focus on value proposition
+
+        4. Purchase Process
+           - Clear explanation of payment process
+           - Generate payment link when customer is ready [https://www.example.com/payment]
+           - Confirm stock availability before purchase
+           - Update inventory after successful transaction
+
+        Communication Style:
+        - Professional yet friendly tone
+        - Clear and concise explanations
+        - Empathetic to customer needs
+        - Use simple, non-technical language
+        - Maintain consistent branding
+        - Limited emoji use (max 1-2 per message)
+
+        Response Templates:
+
+        For out-of-stock items:
+        "I apologize, but [product] is currently out of stock. However, I can recommend [alternative product] which has similar features and is available now. Would you like to learn more about it?"
+
+        For low stock items:
+        "Just to let you know, we only have [X] units of [product] remaining in stock. I'd recommend making your decision soon if you're interested."
+
+        For unrelated products:
+        "I apologize, but as a {company_settings[2]}, we don't carry that item. However, I can suggest [relevant alternative] from our current inventory that might meet your needs. Would you like to learn more?"
+
+        For payment process:
+        "I'm glad you're interested in [product]. I can generate a secure payment link for you to complete your purchase. The current price is $[price]. Would you like to proceed with the purchase of [product]?"
+
+        Remember to:
+        1. Always verify stock before recommending products
+        2. Keep track of customer preferences
+        3. Provide accurate product information
+        4. Update inventory after successful purchases
+        5. Maintain professional communication
         """
 
-        # --- Agent Setup ---
-        prompt_template = ChatPromptTemplate.from_messages([
+        prompt = ChatPromptTemplate.from_messages([
             ("system", system_prompt),
             MessagesPlaceholder(variable_name="chat_history"),
             ("human", "{input}"),
-            MessagesPlaceholder(variable_name="agent_scratchpad") # For tool usage intermediate steps
+            MessagesPlaceholder(variable_name="agent_scratchpad")
         ])
 
-        tools = [product_retriever_tool]
-        agent = create_tool_calling_agent(llm, tools, prompt_template)
-        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True) # Add verbose and error handling
-
-
-        # --- Handle Chat Input ---
-        if user_input := st.chat_input(f"Ask {agent_name} about our products..."):
-            # Add user message to session state chat history (for UI display)
-            st.session_state.chat_history.append({"type": "human", "content": user_input})
+        tools = [tool]
+        agent = create_tool_calling_agent(llm, tools, prompt)
+        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
-            # Display user message in chat message container
+        if prompt_input := st.chat_input("Type your message here..."):
             with st.chat_message("user"):
-                st.write(user_input)
-
-            # Prepare chat history for the agent (Langchain format)
-            formatted_history = []
-            # Use the user object's history which is the ground truth from DB
-            for msg in user.chat_history:
-                if msg["type"] == "human":
-                    formatted_history.append(HumanMessage(content=msg["content"]))
-                elif msg["type"] == "ai":
-                    formatted_history.append(AIMessage(content=msg["content"]))
-                # else: ignore potentially malformed history entries
-
-            # Display AI response in chat message container
+                st.write(prompt_input)
+
             with st.chat_message("assistant"):
-                response_container = st.empty() # Placeholder for potential streaming later
-                ai_response_content = "" # Initialize response content
-
-                try:
-                    with st.spinner(f"{agent_name} is thinking..."):
-                        # Invoke the agent
-                        response = agent_executor.invoke({
-                            "input": user_input,
-                            "chat_history": formatted_history
-                            # agent_scratchpad is handled internally by the executor
-                        })
-                        ai_response_content = response.get("output", "Sorry, I couldn't generate a response.")
-                    response_container.write(ai_response_content)
-
-                except Exception as e:
-                    logger.error(f"Error invoking agent executor: {e}", exc_info=True) # Log stack trace
-                    ai_response_content = f"I apologize, I encountered an technical issue processing your request. Please try again shortly. (Error: {e})"
-                    response_container.error(ai_response_content) # Show error in the chat
-
-            # Add AI response to session state chat history (for UI display)
-            st.session_state.chat_history.append({"type": "ai", "content": ai_response_content})
-
-            # Update user's persistent chat history in DB using the User object method
-            # We pass the latest user message and the AI response
-            new_messages_for_db = [
-                {"type": "human", "content": user_input},
-                {"type": "ai", "content": ai_response_content}
-            ]
-            user.update_chat_history(new_messages_for_db)
-
-
-            # --- Purchase Handling Logic ---
-            # Check if the AI response contains the payment link trigger phrase
-            payment_link = "https://www.example.com/payment"
-            if payment_link in ai_response_content:
-                logger.info(f"Payment link detected in response for user {user.username}.")
-                # Attempt to extract the product name reliably
-                # Look for patterns like "purchase for the [product name]:" or similar
-                product_name = None
-                try:
-                    # Try different patterns to extract the name
-                    patterns = [
-                        r"purchase for the (.*?):",
-                        r"purchase of (.*?)\?",
-                        r"interested in (.*?)\." # Less reliable, add more if needed
-                    ]
-                    import re
-                    for pattern in patterns:
-                        match = re.search(pattern, ai_response_content, re.IGNORECASE)
-                        if match:
-                            product_name = match.group(1).strip().rstrip('.') # Clean up extracted name
-                            logger.info(f"Extracted product name '{product_name}' for purchase.")
-                            break # Stop after first match
-                    if not product_name:
-                         logger.warning("Payment link found, but failed to extract product name from response.")
-                         st.warning("Purchase link mentioned, but couldn't confirm the specific product name from the chat context. Stock not updated.")
-
-                except Exception as extraction_error:
-                    logger.error(f"Error extracting product name: {extraction_error}")
-                    st.warning("Error processing product name for purchase. Stock not updated.")
+                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"]
+                st.write(response)
+                
+                # Handle inventory update on purchase
+                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:
+                        # Retrieve current product details
+                        docs = retrieve_query(product_name)
+                        if docs:
+                            product_doc = docs[0]
+                            product_id = product_doc.metadata.get("id")
+                            current_stock = product_doc.metadata.get("stock")
+                            
+                            if product_id and current_stock > 0:
+                                try:
+                                    conn = sqlite3.connect('users.db')
+                                    c = conn.cursor()
+                                    
+                                    # Verify current stock
+                                    c.execute('SELECT stock FROM products WHERE id = ?', (product_id,))
+                                    actual_stock = c.fetchone()[0]
+                                    
+                                    if actual_stock > 0:
+                                        # Update product stock
+                                        c.execute('UPDATE products SET stock = ? WHERE id = ?', 
+                                                (actual_stock - 1, 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()
+                                        
+                                        # Check for low stock
+                                        if (actual_stock - 1) <= 3 and (actual_stock - 1) > 0:
+                                            st.warning(f"Low stock alert: Only {actual_stock - 1} units of {product_name} remaining!")
+                                    else:
+                                        st.error(f"Sorry, {product_name} is now out of stock!")
+                                except Exception as e:
+                                    st.error(f"Error processing purchase: {str(e)}")
+                                    if conn:
+                                        conn.rollback()
+                                finally:
+                                    if conn:
+                                        conn.close()
+                            else:
+                                st.error("Product is out of stock or not found in the inventory.")
+                        else:
+                            st.error("Could not retrieve product details. Please try again.")
+                    else:
+                        st.error("Please specify which product you'd like to purchase.")
 
+            new_messages = [
+                {"type": "human", "content": prompt_input},
+                {"type": "ai", "content": response}
+            ]
+            st.session_state.user.update_chat_history(new_messages)
+            st.session_state.chat_history += new_messages
 
-                # If product name was extracted, proceed with stock update simulation
-                if product_name:
-                    conn = None # Ensure conn is defined for finally block
-                    try:
-                        conn = get_db_conn()
-                        if not conn: raise sqlite3.Error("Failed to get DB connection for purchase.")
-                        c = conn.cursor()
-
-                        # **Crucial Final Check:** Get current stock directly from DB
-                        c.execute('SELECT id, stock FROM products WHERE LOWER(name) = LOWER(?)', (product_name,))
-                        product_data = c.fetchone()
-
-                        if product_data and product_data[1] > 0: # Check if product exists and stock > 0
-                            product_id = product_data[0]
-                            current_stock = product_data[1]
-                            new_stock = current_stock - 1
-
-                            logger.info(f"Attempting to update stock for '{product_name}' (ID: {product_id}) from {current_stock} to {new_stock}.")
-
-                            # Perform updates in a transaction
-                            c.execute('UPDATE products SET stock = ? WHERE id = ?', (new_stock, product_id))
-                            logger.info(f"Product stock updated in DB for ID {product_id}.")
-
-                            # Update user's purchase history (using the User object method is safer)
-                            # This appends the product name and saves the updated list to DB
-                            user.update_products_bought(product_name)
-                            logger.info(f"User {user.username}'s purchase history updated with '{product_name}'.")
-
-                            conn.commit() # Commit transaction
-                            logger.info(f"Purchase transaction committed for product '{product_name}'.")
-
-                            # Clear the vector store cache AFTER successful commit
-                            load_vector_store.clear()
-                            logger.info("Vector store cache cleared due to stock update.")
-
-                            # Inform user via chat (can't write directly to chat history here easily)
-                            # Use st.success/warning/info below the chat area
-                            st.success(f"Purchase link provided for {product_name}! Inventory updated. Remaining stock: {new_stock}.")
-                            if new_stock <= 3 and new_stock > 0:
-                                st.warning(f"Low stock alert: Only {new_stock} units of {product_name} remaining!")
-                            elif new_stock == 0:
-                                st.warning(f"{product_name} is now out of stock.")
-
-                            # Rerun to potentially update sidebar purchase history display
-                            st.rerun()
-
-
-                        elif product_data and product_data[1] <= 0:
-                            logger.warning(f"Product '{product_name}' (ID: {product_data[0]}) is out of stock (DB check). Purchase aborted.")
-                            st.error(f"Sorry, {product_name} just went out of stock before the transaction could be completed. Please choose another item.")
-                        else:
-                            logger.error(f"Product '{product_name}' not found in database during final purchase check.")
-                            st.error(f"Error: Could not find product '{product_name}' to complete the purchase simulation.")
-
-                    except sqlite3.Error as db_err:
-                        logger.error(f"Database error during purchase update for '{product_name}': {db_err}", exc_info=True)
-                        st.error(f"A database error occurred while processing the purchase for {product_name}. Please try again.")
-                        if conn: conn.rollback() # Rollback on error
-                    except Exception as final_err:
-                        logger.error(f"Unexpected error during purchase update for '{product_name}': {final_err}", exc_info=True)
-                        st.error(f"An unexpected error occurred processing the purchase for {product_name}.")
-                        if conn: conn.rollback()
-                    finally:
-                        if conn: conn.close() # Ensure connection is closed
-
-
-# --- Run the App ---
 if __name__ == "__main__":
-    main()
-
+    main()