Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 import streamlit as st
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
@@ -8,76 +9,111 @@ from langchain.memory import ConversationBufferMemory
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import PyMuPDFLoader
 
-# ──────────────────── 1. Setup ────────────────────
+# App config
+st.set_page_config(page_title="Café Eleven", page_icon="☕")
 st.title("☕ Café Eleven Ordering Assistant")
-st.caption("Powered by LangChain & Streamlit")
 
-# Load documents
-@st.cache_resource
-def load_docs():
-    loader = PyMuPDFLoader("menu.pdf")  # Change to your PDF filename
-    return loader.load()
+# Initialize chat
+if "messages" not in st.session_state:
+    st.session_state.messages = [{
+        "role": "assistant", 
+        "content": "Hi! Welcome to Café Eleven. What would you like to order today?"
+    }]
 
-# Initialize components
+# Display messages
+for msg in st.session_state.messages:
+    st.chat_message(msg["role"]).write(msg["content"])
+
+# Chat functions
 @st.cache_resource
-def init_chain():
-    # Embeddings
+def load_chain():
+    # Load all PDFs in directory
+    pdf_files = [str(p) for p in Path(".").glob("*.pdf")]
+    if not pdf_files:
+        st.error("No PDF files found! Please upload menu PDFs.")
+        st.stop()
+    
+    # Load and split documents
+    docs = []
+    for pdf in pdf_files:
+        loader = PyMuPDFLoader(pdf)
+        docs.extend(loader.load())
+    
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200
+    )
+    splits = text_splitter.split_documents(docs)
+    
+    # Setup vectorstore
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-mpnet-base-v2"
     )
+    vectorstore = FAISS.from_documents(splits, embeddings)
     
-    # Text splitting
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    documents = text_splitter.split_documents(load_docs())
-    
-    # Vectorstore
-    vectorstore = FAISS.from_documents(documents, embeddings)
-    
-    # LLM (using free inference API)
+    # Setup LLM
     llm = HuggingFaceHub(
         repo_id="meta-llama/Llama-2-7b-chat-hf",
         huggingfacehub_api_token=os.environ["HF_TOKEN"],
         model_kwargs={
             "temperature": 0.2,
-            "max_new_tokens": 256
+            "max_length": 256
         }
     )
     
-    # Memory
-    memory = ConversationBufferMemory(
-        memory_key="chat_history",
-        return_messages=True,
-        output_key='answer'
-    )
-    
-    # Chain
+    # Create chain with system prompt
     return ConversationalRetrievalChain.from_llm(
         llm=llm,
         retriever=vectorstore.as_retriever(),
-        memory=memory,
-        return_source_documents=True
+        memory=ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True
+        ),
+        condense_question_prompt="""
+        You are a friendly café assistant. Your job is to:
+        1. Greet customers warmly
+        2. Help them place orders
+        3. Suggest menu items
+        4. Never make up items not in the menu
+        Current conversation:
+        {chat_history}
+        Question: {question}
+        """
     )
 
-# ──────────────────── 2. Chat Interface ────────────────────
-if "messages" not in st.session_state:
-    st.session_state.messages = [
-        {"role": "assistant", "content": "Hi! Welcome to Café Eleven. What would you like to order today?"}
-    ]
-
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
+# Chat input
 if prompt := st.chat_input("Your order..."):
     st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.markdown(prompt)
+    st.chat_message("user").write(prompt)
     
-    with st.chat_message("assistant"):
-        chain = init_chain()
-        result = chain({"question": prompt})
-        response = result["answer"]
-        
-        # Display response
-        st.markdown(response)
-        st.session_state.messages.append({"role": "assistant", "content": response})
+    with st.spinner("Preparing your order..."):
+        try:
+            chain = load_chain()
+            response = chain({"question": prompt})["answer"]
+            
+            # Format response cleanly
+            if "menu" in prompt.lower():
+                response = "Here are our offerings:\n" + response
+            elif "thank" in prompt.lower():
+                response = "You're welcome! " + response
+            
+            st.session_state.messages.append({"role": "assistant", "content": response})
+            st.chat_message("assistant").write(response)
+            
+        except Exception as e:
+            st.error(f"Sorry, something went wrong: {str(e)}")
+
+# PDF upload section
+with st.sidebar:
+    st.subheader("Menu Management")
+    uploaded_files = st.file_uploader(
+        "Upload PDF menus",
+        type="pdf",
+        accept_multiple_files=True
+    )
+    if uploaded_files:
+        for file in uploaded_files:
+            with open(file.name, "wb") as f:
+                f.write(file.getbuffer())
+        st.success(f"Uploaded {len(uploaded_files)} menu(s)")
+        st.cache_resource.clear()  # Refresh the vectorstore