Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,6 +2,7 @@ import os
 import re
 import logging
 from uuid import uuid4
+from pathlib import Path
 from dotenv import load_dotenv
 import streamlit as st
 
@@ -19,6 +20,25 @@ from langchain_chroma import Chroma
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Set up proper cache directories for HuggingFace Spaces
+def setup_environment():
+    # Create cache directories in a writable location
+    cache_dir = Path("/tmp/cache")  # Using /tmp which is writable in HuggingFace Spaces
+    cache_dir.mkdir(exist_ok=True)
+    
+    # Set environment variables
+    os.environ['STREAMLIT_HOME'] = str(cache_dir / "streamlit")
+    os.environ['HF_HOME'] = str(cache_dir / "huggingface")
+    os.environ['TRANSFORMERS_CACHE'] = str(cache_dir / "transformers")
+    os.environ['XDG_CACHE_HOME'] = str(cache_dir)
+    
+    # Ensure subdirectories exist
+    (cache_dir / "huggingface").mkdir(exist_ok=True)
+    (cache_dir / "streamlit").mkdir(exist_ok=True)
+    (cache_dir / "transformers").mkdir(exist_ok=True)
+
+setup_environment()
+
 # Load environment variables
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -30,12 +50,20 @@ if not all([HF_TOKEN, GROQ_API_KEY, PDF_PATH]):
     st.error("Missing required environment variables")
     st.stop()
 
-# Initialize RAG components
-embeddings = HuggingFaceEmbeddings(
-    model_name="sentence-transformers/all-MiniLM-L6-v2",
-    cache_folder="./cache/huggingface"
-)
-llm = ChatGroq(model_name="Deepseek-R1-Distill-Llama-70b")
+# Initialize RAG components with proper cache handling
+try:
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2",
+        model_kwargs={'device': 'cpu'},
+        encode_kwargs={'normalize_embeddings': True},
+        cache_folder=os.environ['HF_HOME']
+    )
+except Exception as e:
+    logger.error(f"Failed to initialize embeddings: {str(e)}")
+    st.error("Failed to initialize embeddings. Please try again later.")
+    st.stop()
+
+llm = ChatGroq(model_name="Deepseek-R1-Distill-Llama-70b", temperature=0.1)
 session_store = {}
 
 # Process PDF into vectorstore
@@ -45,10 +73,13 @@ def process_pdf(file_path: str):
         documents = loader.load()
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=500)
         splits = text_splitter.split_documents(documents)
+        
+        # Use temporary directory for Chroma DB
+        chroma_dir = "/tmp/chroma_db"
         vectorstore = Chroma.from_documents(
             documents=splits,
             embedding=embeddings,
-            persist_directory="./portfolio.db"
+            persist_directory=chroma_dir
         )
         logger.info(f"PDF {file_path} processed successfully")
         return vectorstore
@@ -58,30 +89,35 @@ def process_pdf(file_path: str):
         st.stop()
 
 # Initialize vectorstore and retriever
-vectorstore = process_pdf(PDF_PATH)
-retriever = vectorstore.as_retriever()
+try:
+    vectorstore = process_pdf(PDF_PATH)
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+except Exception as e:
+    logger.error(f"Failed to initialize vectorstore: {str(e)}")
+    st.error("Failed to initialize document store. Please try again later.")
+    st.stop()
 
 # System prompt for the assistant
 system_prompt = """You are Max, a friendly and professional chatbot designed to 
-assist visitors to Nivakaran’s portfolio website. Your primary goal 
+assist visitors to Nivakaran's portfolio website. Your primary goal 
 is to provide accurate, clear, and helpful information about Nivakaran, based 
 on the following context:
 
 {context}
 
 Your responses should be:
-1. Informative and relevant, directly addressing the visitor’s questions about Nivakaran’s skills, 
+1. Informative and relevant, directly addressing the visitor's questions about Nivakaran's skills, 
 projects, experience, and background.
-2. Concise but thorough enough to give visitors a clear understanding of Nivakaran’s expertise.
+2. Concise but thorough enough to give visitors a clear understanding of Nivakaran's expertise.
 3. Engaging and approachable, maintaining a professional yet conversational tone.
-4. Honest about what is available in the provided context; if you don’t know an answer, politely 
+4. Honest about what is available in the provided context; if you don't know an answer, politely 
 say so and suggest the visitor explore other sections of the portfolio or contact Nivakaran directly.
-5. Focused on helping visitors understand Nivakaran’s capabilities and what makes him stand out 
+5. Focused on helping visitors understand Nivakaran's capabilities and what makes him stand out 
 as a developer and professional.
 6. Ready to provide examples, explanations, or links to portfolio projects when relevant.
 
 Avoid providing generic or unrelated information. Always tailor your answers to 
-highlight Nivakaran’s strengths and the unique value he brings.
+highlight Nivakaran's strengths and the unique value he brings.
 """
 
 # Streamlit app UI
@@ -111,10 +147,11 @@ if user_input := st.chat_input("Ask me something about Nivakaran..."):
 
         # Contextualize question based on history
         contextualize_q_prompt = ChatPromptTemplate.from_messages([
-            ("system", "Rephrase questions considering chat history."),
+            ("system", "Given a chat history and the latest user question which might reference context in the chat history, formulate a standalone question which can be understood without the chat history. Return just the question and nothing else."),
             MessagesPlaceholder("chat_history"),
             ("human", "{input}")
         ])
+        
         history_aware_retriever = create_history_aware_retriever(
             llm, retriever, contextualize_q_prompt
         )
@@ -125,6 +162,7 @@ if user_input := st.chat_input("Ask me something about Nivakaran..."):
             MessagesPlaceholder("chat_history"),
             ("human", "{input}")
         ])
+        
         question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
         rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
 
@@ -132,10 +170,11 @@ if user_input := st.chat_input("Ask me something about Nivakaran..."):
             "input": user_input,
             "chat_history": last_messages
         })
+        
         raw_answer = result["answer"]
-
-        # Clean out <think>...</think> junk
+        # Clean out <think>...</think> junk and any other unwanted artifacts
         cleaned_answer = re.sub(r"<think>.*?</think>\s*", "", raw_answer, flags=re.DOTALL).strip()
+        cleaned_answer = re.sub(r"<\|.*?\|>", "", cleaned_answer).strip()
 
         with st.chat_message("assistant"):
             st.markdown(cleaned_answer)
@@ -143,4 +182,5 @@ if user_input := st.chat_input("Ask me something about Nivakaran..."):
         st.session_state.history.add_ai_message(cleaned_answer)
 
     except Exception as e:
-        st.error(f"Error: {str(e)}")
+        logger.error(f"Error during RAG processing: {str(e)}")
+        st.error("Sorry, I encountered an error while processing your request. Please try again.")