Update app.py

app.py

@@ -9,107 +9,110 @@ from langchain.memory import ConversationBufferMemory
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import PyMuPDFLoader
 
-# Initialize session state FIRST (before any other Streamlit calls)
-if 'messages' not in st.session_state:
-    st.session_state.messages = []
+# Initialize session state FIRST (before any other Streamlit operations)
+def initialize_session_state():
+    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?"
+        }]
 
-# Then configure the page
-st.set_page_config(page_title="Café Eleven", page_icon="☕")
-st.title("☕ Café Eleven Ordering Assistant")
-
-# Display welcome message if empty
-if not st.session_state.messages:
-    st.session_state.messages.append({
-        "role": "assistant",
-        "content": "Hi! Welcome to Café Eleven. What would you like to order today?"
-    })
-
-# Display chat messages
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.write(message["content"])
-
-# Chat processing functions
-@st.cache_resource
-def setup_chain():
-    # Load all PDFs
-    pdf_files = list(Path(".").glob("*.pdf"))
-    if not pdf_files:
-        st.error("No PDF menus found! Please upload some.")
-        st.stop()
+# Main application function
+def main():
+    initialize_session_state()
     
-    # Process documents
-    docs = []
-    for pdf in pdf_files:
-        loader = PyMuPDFLoader(str(pdf))
-        docs.extend(loader.load())
+    # App configuration
+    st.set_page_config(page_title="Café Eleven", page_icon="☕")
+    st.title("☕ Café Eleven Ordering Assistant")
     
-    # Split text
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200
-    )
-    splits = text_splitter.split_documents(docs)
+    # Display chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
     
-    # Create vectorstore
-    embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-mpnet-base-v2"
-    )
-    vectorstore = FAISS.from_documents(splits, embeddings)
+    # Chat processing functions
+    @st.cache_resource
+    def setup_chain():
+        pdf_files = list(Path(".").glob("*.pdf"))
+        if not pdf_files:
+            st.error("No PDF menus found! Please upload some.")
+            st.stop()
+        
+        # Process all PDFs
+        docs = []
+        for pdf in pdf_files:
+            loader = PyMuPDFLoader(str(pdf))
+            docs.extend(loader.load())
+        
+        # Split text
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200
+        )
+        splits = text_splitter.split_documents(docs)
+        
+        # Create embeddings
+        embeddings = HuggingFaceEmbeddings(
+            model_name="sentence-transformers/all-mpnet-base-v2"
+        )
+        vectorstore = FAISS.from_documents(splits, embeddings)
+        
+        # Setup LLM
+        llm = HuggingFaceHub(
+            repo_id="meta-llama/Llama-2-7b-chat-hf",
+            huggingfacehub_api_token=os.environ.get("HF_TOKEN", ""),
+            model_kwargs={
+                "temperature": 0.2,
+                "max_length": 256
+            }
+        )
+        
+        # Create conversation chain
+        return ConversationalRetrievalChain.from_llm(
+            llm=llm,
+            retriever=vectorstore.as_retriever(),
+            memory=ConversationBufferMemory(
+                memory_key="chat_history",
+                return_messages=True
+            )
+        )
     
-    # Setup LLM
-    llm = HuggingFaceHub(
-        repo_id="meta-llama/Llama-2-7b-chat-hf",
-        huggingfacehub_api_token=os.environ.get("HF_TOKEN", ""),
-        model_kwargs={
-            "temperature": 0.2,
-            "max_length": 256
-        }
-    )
-    
-    # Create conversation chain
-    return ConversationalRetrievalChain.from_llm(
-        llm=llm,
-        retriever=vectorstore.as_retriever(),
-        memory=ConversationBufferMemory(
-            memory_key="chat_history",
-            return_messages=True
-        ),
-        get_chat_history=lambda h: h
-    )
-
-# Chat input handling
-if prompt := st.chat_input("What would you like to order?"):
-    # Add user message
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.write(prompt)
+    # Handle chat input
+    if prompt := st.chat_input("What would you like to order?"):
+        # Add user message
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.write(prompt)
+        
+        # Generate response
+        with st.chat_message("assistant"):
+            with st.spinner("Preparing your order..."):
+                try:
+                    chain = setup_chain()
+                    response = chain({"question": prompt})["answer"]
+                    st.write(response)
+                    st.session_state.messages.append({
+                        "role": "assistant",
+                        "content": response
+                    })
+                except Exception as e:
+                    st.error(f"Error: {str(e)}")
     
-    # Generate response
-    with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            try:
-                chain = setup_chain()
-                response = chain({"question": prompt})["answer"]
-                st.write(response)
-                st.session_state.messages.append({
-                    "role": "assistant",
-                    "content": response
-                })
-            except Exception as e:
-                st.error(f"Error generating response: {str(e)}")
+    # PDF upload in sidebar
+    with st.sidebar:
+        st.header("Menu Management")
+        uploaded_files = st.file_uploader(
+            "Upload menu PDFs",
+            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)} new menu(s)")
+            st.cache_resource.clear()
 
-# PDF upload in sidebar
-with st.sidebar:
-    st.header("Menu Management")
-    uploaded_files = st.file_uploader(
-        "Upload menu PDFs",
-        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)} new menu(s)")
-        st.cache_resource.clear()
+# Entry point for Hugging Face Spaces
+if __name__ == "__main__":
+    main()